diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean index 09b114bbb8..c27d14d07d 100644 --- a/src/Init/Data/Array/Basic.lean +++ b/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. diff --git a/src/Lean/Compiler/LCNF/ExplicitRC.lean b/src/Lean/Compiler/LCNF/ExplicitRC.lean index f90c6d1d3b..785ae0c815 100644 --- a/src/Lean/Compiler/LCNF/ExplicitRC.lean +++ b/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 diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h index 01156d100d..87cab7e9b3 100644 --- a/src/include/lean/lean.h +++ b/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)); } diff --git a/tests/lean/4240.lean.expected.out b/tests/lean/4240.lean.expected.out index aa7e713870..b4c3b1dbef 100644 --- a/tests/lean/4240.lean.expected.out +++ b/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; diff --git a/tests/lean/run/ugetBorrowed.lean b/tests/lean/run/ugetBorrowed.lean new file mode 100644 index 0000000000..c9c01281da --- /dev/null +++ b/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)