Inner recursive worker with memoization keyed by `(sub-expr addr, depth)`. Depth enters the key because traversing under a binder increments `depth`, and the substitution's semantics change: under one extra binder, `Var(depth+1)` now targets the original substitution site. Two subtrees with the same address but visited at different depths must not share a result.
( env: &mut InternTable<M>, body: &KExpr<M>, arg: &KExpr<M>, depth: u64, cache: &mut FxHashMap<(Addr, u64), KExpr<M>>, )
| 138 | /// depth)`. Depth enters the key because traversing under a binder |
| 139 | /// increments `depth`, and the substitution's semantics change: under |
| 140 | /// one extra binder, `Var(depth+1)` now targets the original |
| 141 | /// substitution site. Two subtrees with the same address but visited at |
| 142 | /// different depths must not share a result. |
| 143 | fn subst_cached<M: KernelMode>( |
| 144 | env: &mut InternTable<M>, |
| 145 | body: &KExpr<M>, |
| 146 | arg: &KExpr<M>, |
| 147 | depth: u64, |
| 148 | cache: &mut FxHashMap<(Addr, u64), KExpr<M>>, |
| 149 | ) -> KExpr<M> { |
| 150 | if body.lbr() <= depth { |
| 151 | return body.clone(); |
| 152 | } |
| 153 | |
| 154 | // Pointer-identity cache: expressions are content-addressed, so two |
| 155 | // sub-trees with the same `addr()` are structurally equal, meaning |
| 156 | // `subst` at the same `depth` must produce the same result. Skipping |
| 157 | // re-traversal here is the whole point of the cache — for Lean bodies |
| 158 | // with significant sub-term sharing it turns an O(tree-size) walk |
| 159 | // into O(dag-size). |
| 160 | let key = (body.hash_key(), depth); |
| 161 | if let Some(cached) = cache.get(&key) { |
| 162 | return cached.clone(); |
| 163 | } |
| 164 | |
| 165 | let result = match body.data() { |
| 166 | ExprData::Var(i, name, _) => { |
| 167 | let i = *i; |
| 168 | if i == depth { |
| 169 | lift(env, arg, depth, 0) |
| 170 | } else if i > depth { |
| 171 | KExpr::var(i - 1, name.clone()) |
| 172 | } else { |
| 173 | // Unreachable under the outer `lbr() <= depth` guard (Var below |
| 174 | // `depth` is bound, so its lbr is below depth and we'd have |
| 175 | // returned early), but keep the explicit branch for clarity. |
| 176 | let r = body.clone(); |
| 177 | cache.insert(key, r.clone()); |
| 178 | return r; |
| 179 | } |
| 180 | }, |
| 181 | |
| 182 | ExprData::App(f, x, _) => { |
| 183 | let f2 = subst_cached(env, f, arg, depth, cache); |
| 184 | let x2 = subst_cached(env, x, arg, depth, cache); |
| 185 | KExpr::app(f2, x2) |
| 186 | }, |
| 187 | |
| 188 | ExprData::Lam(name, bi, ty, inner, _) => { |
| 189 | let ty2 = subst_cached(env, ty, arg, depth, cache); |
| 190 | let inner2 = subst_cached(env, inner, arg, depth + 1, cache); |
| 191 | KExpr::lam(name.clone(), bi.clone(), ty2, inner2) |
| 192 | }, |
| 193 | |
| 194 | ExprData::All(name, bi, ty, inner, _) => { |
| 195 | let ty2 = subst_cached(env, ty, arg, depth, cache); |
| 196 | let inner2 = subst_cached(env, inner, arg, depth + 1, cache); |
| 197 | KExpr::all(name.clone(), bi.clone(), ty2, inner2) |