Rewrite an expression tree to use Share(idx) references. Uses iterative traversal with caching to handle deep trees and Arc sharing.
( expr: &Arc<Expr>, hash_to_idx: &HashMap<blake3::Hash, u64>, ptr_to_hash: &FxHashMap<*const Expr, blake3::Hash>, cache: &mut FxHashMap<*const Expr, Arc<Expr>>, )
| 694 | /// Rewrite an expression tree to use Share(idx) references. |
| 695 | /// Uses iterative traversal with caching to handle deep trees and Arc sharing. |
| 696 | fn rewrite_expr( |
| 697 | expr: &Arc<Expr>, |
| 698 | hash_to_idx: &HashMap<blake3::Hash, u64>, |
| 699 | ptr_to_hash: &FxHashMap<*const Expr, blake3::Hash>, |
| 700 | cache: &mut FxHashMap<*const Expr, Arc<Expr>>, |
| 701 | ) -> Arc<Expr> { |
| 702 | let mut stack: Vec<RewriteFrame<'_>> = vec![RewriteFrame::Visit(expr)]; |
| 703 | let mut results: Vec<Arc<Expr>> = Vec::new(); |
| 704 | |
| 705 | while let Some(frame) = stack.pop() { |
| 706 | match frame { |
| 707 | RewriteFrame::Visit(e) => { |
| 708 | let ptr = e.as_ref() as *const Expr; |
| 709 | |
| 710 | // Check hash_to_idx FIRST: if this expression is shareable, replace |
| 711 | // it with Share(idx) even if the cache has a stale (pre-sharing) |
| 712 | // entry. This ordering eliminates the need for cache.clear() in the |
| 713 | // outer build_sharing_vec loop. |
| 714 | if let Some(hash) = ptr_to_hash.get(&ptr) |
| 715 | && let Some(&idx) = hash_to_idx.get(hash) |
| 716 | { |
| 717 | let share = Expr::share(idx); |
| 718 | cache.insert(ptr, share.clone()); |
| 719 | results.push(share); |
| 720 | continue; |
| 721 | } |
| 722 | |
| 723 | // Cache hit for non-shareable sub-expressions |
| 724 | if let Some(cached) = cache.get(&ptr) { |
| 725 | results.push(cached.clone()); |
| 726 | continue; |
| 727 | } |
| 728 | |
| 729 | // Process based on node type |
| 730 | match e.as_ref() { |
| 731 | // Leaf nodes - return as-is |
| 732 | Expr::Sort(_) |
| 733 | | Expr::Var(_) |
| 734 | | Expr::Ref(..) |
| 735 | | Expr::Rec(..) |
| 736 | | Expr::Str(_) |
| 737 | | Expr::Nat(_) |
| 738 | | Expr::Share(_) => { |
| 739 | cache.insert(ptr, e.clone()); |
| 740 | results.push(e.clone()); |
| 741 | }, |
| 742 | |
| 743 | // Nodes with children - push build frame, then visit children |
| 744 | Expr::Prj(type_ref_idx, field_idx, val) => { |
| 745 | stack.push(RewriteFrame::BuildPrj(e, *type_ref_idx, *field_idx)); |
| 746 | stack.push(RewriteFrame::Visit(val)); |
| 747 | }, |
| 748 | Expr::App(fun, arg) => { |
| 749 | stack.push(RewriteFrame::BuildApp(e)); |
| 750 | stack.push(RewriteFrame::Visit(arg)); |
| 751 | stack.push(RewriteFrame::Visit(fun)); |
| 752 | }, |
| 753 | Expr::Lam(ty, body) => { |