From ecce5e69bf6dabbfad6283cb8844d17596320b06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20B=C3=B6ving?= <hargonix@gmail.com>
Date: Mon, 8 Dec 2025 11:55:46 +0100
Subject: [PATCH] feat: tagged_return attribute (#11530)

This PR introduces the new `tagged_return` attribute. It allows users to
mark `extern` declarations to be guaranteed to always return `tagged`
return values. Unlike with `object` or `tobject` the compiler does not
emit reference counting operations for them. In the future information
from this attribute will be used for a more powerful analysis to remove
reference counts when possible.
---
 src/Lean/Compiler/IR/ToIR.lean      | 18 ++++++++++++++-
 src/stdlib_flags.h                  |  2 ++
 tests/lean/run/tagged_return_1.lean | 36 +++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)
 create mode 100644 tests/lean/run/tagged_return_1.lean

diff --git a/src/Lean/Compiler/IR/ToIR.lean b/src/Lean/Compiler/IR/ToIR.lean
index 75ca86f7a4..a3ec28a2ce 100644
--- a/src/Lean/Compiler/IR/ToIR.lean
+++ b/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.
diff --git a/src/stdlib_flags.h b/src/stdlib_flags.h
index 79a0e58edd..6c6b3bee8f 100644
--- a/src/stdlib_flags.h
+++ b/src/stdlib_flags.h
@@ -1,5 +1,7 @@
 #include "util/options.h"
 
+// update thy
+
 namespace lean {
 options get_default_options() {
     options opts;
diff --git a/tests/lean/run/tagged_return_1.lean b/tests/lean/run/tagged_return_1.lean
new file mode 100644
index 0000000000..81a6401c11
--- /dev/null
+++ b/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