feat: add isExclusiveUnsafe

Avoid potentially expensive `e.replace` if it is not applicable.

src/Init/Util.lean

@@ -45,6 +45,13 @@ def dbgSleep {α : Type u} (ms : UInt32) (f : Unit → α) : α := f ()
 @[extern "lean_ptr_addr"]
 unsafe opaque ptrAddrUnsafe {α : Type u} (a : @& α) : USize
 
+/--
+Returns `true` if `a` is an exclusive object.
+We say an object is exclusive if it is single-threaded and its reference counter is 1.
+-/
+@[extern "lean_is_exclusive_obj"]
+unsafe opaque isExclusiveUnsafe {α : Type u} (a : @& α) : Bool
+
 set_option linter.unusedVariables.funArgs false in
 @[inline] unsafe def withPtrAddrUnsafe {α : Type u} {β : Type v} (a : α) (k : USize → β) (h : ∀ u₁ u₂, k u₁ = k u₂) : β :=
   k (ptrAddrUnsafe a)

src/Lean/Data/RBMap.lean

@@ -80,6 +80,10 @@ protected def max : RBNode α β → Option (Sigma (fun k => β k))
 def singleton (k : α) (v : β k) : RBNode α β :=
   node red leaf k v leaf
 
+def isSingleton : RBNode α β → Bool
+  | node _ leaf _ _ leaf => true
+  | _ => false
+
 -- the first half of Okasaki's `balance`, concerning red-red sequences in the left child
 @[inline] def balance1 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
   | node red (node red a kx vx b) ky vy c, kz, vz, d
@@ -269,6 +273,9 @@ variable {α : Type u} {β : Type v} {σ : Type w} {cmp : α → α → Ordering
 def depth (f : Nat → Nat → Nat) (t : RBMap α β cmp) : Nat :=
   t.val.depth f
 
+def isSingleton (t : RBMap α β cmp) : Bool :=
+  t.val.isSingleton
+
 @[inline] def fold (f : σ → α → β → σ) : (init : σ) → RBMap α β cmp → σ
   | b, ⟨t, _⟩ => t.fold f b
 

src/Lean/Elab/MutualDef.lean

@@ -728,12 +728,26 @@ def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecC
   letRecClosures.foldl (init := r) fun r c =>
     r.insert c.toLift.fvarId c.closed
 
+def isApplicable (r : Replacement) (e : Expr) : Bool :=
+  Option.isSome <| e.findExt? fun e =>
+    if e.hasFVar then
+      match e with
+      | .fvar fvarId => if r.contains fvarId then .found else .done
+      | _ => .visit
+    else
+      .done
+
 def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
-  e.replace fun e => match e with
-    | .fvar fvarId => match r.find? fvarId with
-      | some c => some c
-      | _      => none
-    | _ => none
+  -- Remark: if `r` is not a singlenton, then declaration is using `mutual` or `let rec`,
+  -- and there is a big chance `isApplicable r e` is true.
+  if r.isSingleton && !isApplicable r e then
+    e
+  else
+    e.replace fun e => match e with
+      | .fvar fvarId => match r.find? fvarId with
+        | some c => some c
+        | _      => none
+      | _ => none
 
 def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
     : TermElabM (Array PreDefinition) :=

src/include/lean/lean.h

@@ -481,6 +481,10 @@ static inline bool lean_is_exclusive(lean_object * o) {
     }
 }
 
+static inline uint8_t lean_is_exclusive_obj(lean_object * o) {
+    return lean_is_exclusive(o);
+}
+
 static inline bool lean_is_shared(lean_object * o) {
     if (LEAN_LIKELY(lean_is_st(o))) {
         return o->m_rc > 1;
@@ -1133,6 +1137,17 @@ static inline void * lean_get_external_data(lean_object * o) {
     return lean_to_external(o)->m_data;
 }
 
+static inline lean_object * lean_set_external_data(lean_object * o, void * data) {
+    if (lean_is_exclusive(o)) {
+        lean_to_external(o)->m_data = data;
+        return o;
+    } else {
+        lean_object * o_new = lean_alloc_external(lean_get_external_class(o), data);
+        lean_dec_ref(o);
+        return o_new;
+    }
+}
+
 /* Natural numbers */
 
 #define LEAN_MAX_SMALL_NAT (SIZE_MAX >> 1)

src/runtime/sharecommon.cpp

@@ -358,30 +358,16 @@ class sharecommon_quick_fn {
         return result;
     }
 
-    lean_object * visit_sarray(lean_object * a) {
-        lean_object * r = check_cache(a);
-        if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
-        size_t sz        = lean_sarray_size(a);
-        unsigned elem_sz = lean_sarray_elem_size(a);
-        lean_sarray_object * new_a = (lean_sarray_object*)lean_alloc_sarray(elem_sz, sz, sz);
-        memcpy(new_a->m_data, lean_to_sarray(a)->m_data, elem_sz*sz);
-        return save(a, (lean_object*)new_a);
-    }
-
-    lean_object * visit_string(lean_object * a) {
-        lean_object * r = check_cache(a);
-        if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
-        size_t sz     = lean_string_size(a);
-        size_t len    = lean_string_len(a);
-        lean_string_object * new_a = (lean_string_object*)lean_alloc_string(sz, sz, len);
-        lean_set_st_header((lean_object*)new_a, LeanString, 0);
-        new_a->m_size     = sz;
-        new_a->m_capacity = sz;
-        new_a->m_length   = len;
-        memcpy(new_a->m_data, lean_to_string(a)->m_data, sz);
-        return save(a, (lean_object*)new_a);
+    // `sarray` and `string`
+    lean_object * visit_terminal(lean_object * a) {
+        auto it = m_set.find(a);
+        if (it == m_set.end()) {
+            m_set.insert(a);
+        } else {
+            a = *it;
+        }
+        lean_inc_ref(a);
+        return a;
     }
 
     lean_object * visit_array(lean_object * a) {
@@ -399,7 +385,6 @@ class sharecommon_quick_fn {
     lean_object * visit_ctor(lean_object * a) {
         lean_object * r = check_cache(a);
         if (r != nullptr) { lean_assert(r->m_rc > 1); return r; }
-
         unsigned num_objs      = lean_ctor_num_objs(a);
         unsigned tag           = lean_ptr_tag(a);
         unsigned sz            = lean_object_byte_size(a);
@@ -442,9 +427,9 @@ public:
         case LeanRef:             lean_inc_ref(a); return a;
         case LeanExternal:        lean_inc_ref(a); return a;
         case LeanReserved:        lean_inc_ref(a); return a;
+        case LeanScalarArray:     return visit_terminal(a);
+        case LeanString:          return visit_terminal(a);
         case LeanArray:           return visit_array(a);
-        case LeanScalarArray:     return visit_sarray(a);
-        case LeanString:          return visit_string(a);
         default:                  return visit_ctor(a);
         }
     }