feat: derived value analysis for Array.uget (#12604)

This PR makes the derived value analysis in RC insertion recognize
`Array.uget` as another kind of
"projection-like" operation. This allows it to reduce reference count
pressure on elements accessed
through uget.

src/Init/Data/Array/Basic.lean

@@ -170,6 +170,15 @@ This avoids overhead due to unboxing a `Nat` used as an index.
 def uget (xs : @& Array α) (i : USize) (h : i.toNat < xs.size) : α :=
   xs[i.toNat]
 
+/--
+Version of `Array.uget` that does not increment the reference count of its result.
+
+This is only intended for direct use by the compiler.
+-/
+@[extern "lean_array_uget_borrowed"]
+unsafe opaque ugetBorrowed (xs : @& Array α) (i : USize) (h : i.toNat < xs.size) : α :=
+  xs.uget i h
+
 /--
 Low-level modification operator which is as fast as a C array write. The modification is performed
 in-place when the reference to the array is unique.

src/Lean/Compiler/LCNF/ExplicitRC.lean

@@ -104,6 +104,9 @@ partial def collectCode (code : Code .impure) : M Unit := do
     | .fap ``Array.get!Internal args =>
       if let .fvar parent := args[2]! then
         addDerivedValue parent decl.fvarId
+    | .fap ``Array.uget args =>
+      if let .fvar parent := args[1]! then
+        addDerivedValue parent decl.fvarId
     | .reset _ target =>
       removeFromParent target
     | _ => pure ()
@@ -556,6 +559,8 @@ def LetDecl.explicitRc (code : Code .impure) (decl : LetDecl .impure) (k : Code
           pure <| .fap ``Array.getInternalBorrowed args
         else if f == ``Array.get!Internal && (← isBorrowed decl.fvarId) then
           pure <| .fap ``Array.get!InternalBorrowed args
+        else if f == ``Array.uget && (← isBorrowed decl.fvarId) then
+          pure <| .fap ``Array.ugetBorrowed args
         else
           pure <| decl.value
       let decl ← decl.updateValue value

src/include/lean/lean.h

@@ -826,6 +826,10 @@ static inline lean_object * lean_array_uget(b_lean_obj_arg a, size_t i) {
     return r;
 }
 
+static inline b_lean_obj_res lean_array_uget_borrowed(b_lean_obj_arg a, size_t i) {
+    return lean_array_get_core(a, i);
+}
+
 static inline lean_obj_res lean_array_fget(b_lean_obj_arg a, b_lean_obj_arg i) {
     return lean_array_uget(a, lean_unbox(i));
 }

tests/lean/4240.lean.expected.out

@@ -15,9 +15,8 @@
 [Compiler.IR] [result]
     def isSomeWithInstanceNat (x_1 : @& obj) : u8 :=
       let x_2 : usize := 0;
-      let x_3 : tobj := Array.uget ◾ x_1 x_2 ◾;
+      let x_3 : tobj := Array.ugetBorrowed ◾ x_1 x_2 ◾;
       let x_4 : u8 := MyOption.isSomeWithInstance._at_.isSomeWithInstanceNat.spec_0 x_3;
-      dec x_3;
       ret x_4
     def isSomeWithInstanceNat._boxed (x_1 : obj) : tagged :=
       let x_2 : u8 := isSomeWithInstanceNat x_1;

tests/lean/run/ugetBorrowed.lean

@@ -0,0 +1,36 @@
+set_option trace.Compiler.explicitRc true in
+/--
+trace: [Compiler.explicitRc] size: 3
+    def test._redArg @&xs i : UInt8 :=
+      let _x.1 := 5;
+      let _x.2 := Array.ugetBorrowed ◾ xs i ◾;
+      let _x.3 := Nat.decLt _x.1 _x.2;
+      return _x.3
+[Compiler.explicitRc] size: 5
+    def test._redArg._boxed xs i : tagged :=
+      let i.boxed := unbox i;
+      dec i;
+      let res := test._redArg xs i.boxed;
+      dec xs;
+      let r := box res;
+      return r
+[Compiler.explicitRc] size: 1
+    def test @&xs i h : UInt8 :=
+      let _x.1 := test._redArg xs i;
+      return _x.1
+[Compiler.explicitRc] size: 5
+    def test._boxed xs i h : tagged :=
+      let i.boxed := unbox i;
+      dec i;
+      let res := test xs i.boxed h;
+      dec xs;
+      let r := box res;
+      return r
+-/
+#guard_msgs in
+@[noinline] def test (xs : @& Array Nat) (i : USize) (h : i.toNat < xs.size) : Bool :=
+  xs.uget i h > 5
+
+/-- info: true -/
+#guard_msgs in
+#eval test #[1, 2, 10] 2 (by native_decide)