feat: tagged_return attribute

src/Lean/Compiler/IR/ToIR.lean

@@ -19,6 +19,14 @@ open Lean.Compiler (LCNF.Alt LCNF.Arg LCNF.Code LCNF.Decl LCNF.DeclValue LCNF.LC
 
 namespace ToIR
 
+/--
+Marks an extern definition to be guaranteed to always return tagged values.
+This information is used to optimize reference counting in the compiler.
+-/
+@[builtin_doc]
+builtin_initialize taggedReturnAttr : TagAttribute ←
+  registerTagAttribute `tagged_return "mark extern definition to always return tagged values"
+
 structure BuilderState where
   vars : Std.HashMap FVarId Arg := {}
   joinPoints : Std.HashMap FVarId JoinPointId := {}
@@ -334,12 +342,20 @@ where resultTypeForArity (type : Lean.Expr) (arity : Nat) : Lean.Expr :=
 
 def lowerDecl (d : LCNF.Decl) : M (Option Decl) := do
   let params ← d.params.mapM lowerParam
-  let resultType ← lowerResultType d.type d.params.size
+  let mut resultType ← lowerResultType d.type d.params.size
+  let taggedReturn := taggedReturnAttr.hasTag (← getEnv) d.name
   match d.value with
   | .code code =>
+    if taggedReturn then
+      throwError m!"Error while compiling function '{d.name}': @[tagged_return] is only valid for extern declarations"
     let body ← lowerCode code
     pure <| some <| .fdecl d.name params resultType body {}
   | .extern externAttrData =>
+    if taggedReturn then
+      if resultType.isScalar then
+        throwError m!"@[tagged_return] on function '{d.name}' with scalar return type {resultType}"
+      else
+        resultType := .tagged
     if externAttrData.entries.isEmpty then
       -- TODO: This matches the behavior of the old compiler, but we should
       -- find a better way to handle this.

src/stdlib_flags.h

@@ -1,5 +1,7 @@
 #include "util/options.h"
 
+// update thy
+
 namespace lean {
 options get_default_options() {
     options opts;

tests/lean/run/tagged_return_1.lean

@@ -0,0 +1,36 @@
+/-! This test asserts the basic behavior of the tagged_return attribute -/
+
+@[extern "mytest", tagged_return]
+opaque test (a : Nat) : Nat
+
+/--
+trace: [Compiler.IR] [result]
+    def useTest (x_1 : tobj) (x_2 : @& tobj) : tobj :=
+      inc x_1;
+      let x_3 : tagged := test x_1;
+      let x_4 : tobj := Nat.add x_3 x_1;
+      dec x_1;
+      let x_5 : tobj := Nat.add x_4 x_2;
+      dec x_4;
+      ret x_5
+    def useTest._boxed (x_1 : tobj) (x_2 : tobj) : tobj :=
+      let x_3 : tobj := useTest x_1 x_2;
+      dec x_2;
+      ret x_3
+-/
+#guard_msgs in
+set_option trace.compiler.ir.result true in
+def useTest (a b : Nat) :=
+  (test a + a) + b
+
+/--
+error: Error while compiling function 'illegal1': @[tagged_return] is only valid for extern declarations
+-/
+#guard_msgs in
+@[tagged_return]
+opaque illegal1 (a : Nat) : Nat
+
+/-- error: @[tagged_return] on function 'illegal2' with scalar return type u8 -/
+#guard_msgs in
+@[extern "mytest", tagged_return]
+opaque illegal2 (a : Nat) : UInt8