diff --git a/src/compression.rs b/src/compression.rs index 3625ee8e..26285ea4 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -312,7 +312,7 @@ impl Hash for Pattern { } /// only used during tracking - gets the zippers to args of a pattern -fn zids_of_ivar_of_expr(expr: &ExprOwned, zid_of_zip: &FxHashMap) -> Option>> { +fn zids_of_ivar_of_expr(expr: &ExprOwned, zippers: &Zippers) -> Option>> { // quickly determine arity let mut arity = 0; @@ -327,31 +327,34 @@ fn zids_of_ivar_of_expr(expr: &ExprOwned, zid_of_zip: &FxHashMap) -> let mut curr_zip: Zipper = Zipper::default(); let mut zids_of_ivar = vec![vec![]; arity as usize]; - fn helper(expr: Expr, curr_zip: &mut Zipper, zids_of_ivar: &mut Vec>, zid_of_zip: &FxHashMap) -> Result<(), ()> { + fn helper(expr: Expr, curr_zip: &mut Zipper, zids_of_ivar: &mut Vec>, zippers: &Zippers) -> Result<(), ()> { match expr.node() { Node::Prim(_) => {}, Node::Var(_, _) => {}, Node::IVar(i) => { - zids_of_ivar[*i as usize].push(zid_of_zip.get(curr_zip).cloned().ok_or(())?); + let Some(zid) = zippers.get_interned_idx(curr_zip) else { + return Err(()); + }; + zids_of_ivar[*i as usize].push(zid); }, Node::Lam(b, _) => { curr_zip.add_to_end(ZNode::Body); - helper(expr.get(*b), curr_zip, zids_of_ivar, zid_of_zip)?; + helper(expr.get(*b), curr_zip, zids_of_ivar, zippers)?; curr_zip.remove_from_end(); } Node::App(f,x) => { curr_zip.add_to_end(ZNode::Func); - helper(expr.get(*f), curr_zip, zids_of_ivar, zid_of_zip)?; + helper(expr.get(*f), curr_zip, zids_of_ivar, zippers)?; curr_zip.remove_from_end(); curr_zip.add_to_end(ZNode::Arg); - helper(expr.get(*x), curr_zip, zids_of_ivar, zid_of_zip)?; + helper(expr.get(*x), curr_zip, zids_of_ivar, zippers)?; curr_zip.remove_from_end(); } } Ok(()) } // we can pick any match location - if helper(expr.immut(), &mut curr_zip, &mut zids_of_ivar, zid_of_zip).is_err() { + if helper(expr.immut(), &mut curr_zip, &mut zids_of_ivar, zippers).is_err() { return None }; @@ -510,9 +513,9 @@ impl Pattern { let mut curr_zip: Zipper = Zipper::default(); // map zids to zips with a bool thats true if this is a hole and false if its a future ivar - let zips: Vec<(Zipper,Node)> = self.holes.iter().map(|zid| (shared.zip_of_zid[*zid].clone(), Node::Prim(HOLE_SYM.clone()))) + let zips: Vec<(Zipper,Node)> = self.holes.iter().map(|zid| (shared.zippers.zip_of_zid[*zid].clone(), Node::Prim(HOLE_SYM.clone()))) .chain(self.pattern_args.iterate_arguments() - .map(|labelled_zid| (shared.zip_of_zid[labelled_zid.zid].clone(), Node::IVar(labelled_zid.ivar as i32)))).collect(); + .map(|labelled_zid| (shared.zippers.zip_of_zid[labelled_zid.zid].clone(), Node::IVar(labelled_zid.ivar as i32)))).collect(); fn helper(set: &mut ExprSet, curr_node: Idx, curr_zip: &mut Zipper, zips: &[(Zipper,Node)], shared: &SharedData) -> Idx { if let Some((_,e)) = zips.iter().find(|(zip,_)| zip == curr_zip) { @@ -550,7 +553,7 @@ impl Pattern { let mut expr = self.to_expr(shared); let expands_to = format!("{}",tracked_expands_to(self, hole_zid, shared)).magenta().bold().to_string(); let replace_sentinel = Node::Prim("".into()); - let idx = shared.zip_of_zid[hole_zid].zip(&expr); + let idx = shared.zippers.zip_of_zid[hole_zid].zip(&expr); expr.set[idx] = replace_sentinel; expr.to_string().replace("", &expands_to) } @@ -617,7 +620,7 @@ pub struct CriticalMultithreadData { pub struct SharedData { pub crit: Mutex, pub programs: Vec, - pub arg_of_zid_node: Vec>, + pub zippers: Zippers, pub cost_fn: ExprCost, pub analyzed_free_vars: AnalyzedExpr, pub sym_var_info: Option, @@ -626,8 +629,6 @@ pub struct SharedData { pub corpus_span: Span, pub roots: Vec, pub zids_of_node: FxHashMap>, - pub zip_of_zid: Vec, - pub zid_of_zip: FxHashMap, pub extensions_of_zid: Vec, pub set: ExprSet, pub num_paths_to_node: Vec, @@ -783,17 +784,17 @@ impl HoleChoice { }, HoleChoice::FewApps => { pattern.holes.iter().enumerate().map(|(hole_idx,hole_zid)| - (hole_idx, pattern.match_locations.iter().filter(|loc|shared.arg_of_zid_node[*hole_zid][loc].expands_to.is_app()).count())) + (hole_idx, pattern.match_locations.iter().filter(|loc|shared.zippers.arg_of_zid_node[*hole_zid][loc].expands_to.is_app()).count())) .min_by_key(|x|x.1).unwrap().0 } HoleChoice::MaxCost => { pattern.holes.iter().enumerate().map(|(hole_idx,hole_zid)| - (hole_idx, pattern.match_locations.iter().map(|loc|shared.arg_of_zid_node[*hole_zid][loc].cost).sum::())) + (hole_idx, pattern.match_locations.iter().map(|loc|shared.zippers.arg_of_zid_node[*hole_zid][loc].cost).sum::())) .max_by_key(|x|x.1).unwrap().0 } HoleChoice::MinCost => { pattern.holes.iter().enumerate().map(|(hole_idx,hole_zid)| - (hole_idx, pattern.match_locations.iter().map(|loc|shared.arg_of_zid_node[*hole_zid][loc].cost).sum::())) + (hole_idx, pattern.match_locations.iter().map(|loc|shared.zippers.arg_of_zid_node[*hole_zid][loc].cost).sum::())) .min_by_key(|x|x.1).unwrap().0 } HoleChoice::MaxLargestSubset => { @@ -801,7 +802,7 @@ impl HoleChoice { // mainly because where there are like dozens of holes doing all these lookups and clones and hashmaps is a LOT pattern.holes.iter().enumerate() .map(|(hole_idx,hole_zid)| (hole_idx, *pattern.match_locations.iter() - .map(|loc| shared.arg_of_zid_node[*hole_zid][loc].expands_to.clone()).counts().values().max().unwrap())).max_by_key(|&(_,max_count)| max_count).unwrap().0 + .map(|loc| shared.zippers.arg_of_zid_node[*hole_zid][loc].expands_to.clone()).counts().values().max().unwrap())).max_by_key(|&(_,max_count)| max_count).unwrap().0 } _ => unimplemented!() } @@ -935,7 +936,7 @@ fn stitch_search( // get the hashmap for looking up the Arg struct for this hole based on match location. The Arg // struct has a bunch of info about the hole, including what it expands into at each match location - let arg_of_loc = &shared.arg_of_zid_node[hole_zid]; + let arg_of_loc = &shared.zippers.arg_of_zid_node[hole_zid]; // sort the match locations by node type (ie what theyll expand into) so that we can do a group_by() on // node type in order to iterate over all the different expansions @@ -965,13 +966,13 @@ fn stitch_search( if should_prune_single_task(&shared, &locs) { if !shared.cfg.no_stats { shared.stats.lock().deref_mut().single_task_fired += 1; } - if tracked && !shared.cfg.quiet { println!("{} single task pruned when expanding {} to {}", "[TRACK]".red().bold(), original_pattern.to_expr(&shared), zipper_replace(original_pattern.to_expr(&shared), &shared.zip_of_zid[hole_zid], Node::Prim(format!("<{expands_to}>").into()))) } + if tracked && !shared.cfg.quiet { println!("{} single task pruned when expanding {} to {}", "[TRACK]".red().bold(), original_pattern.to_expr(&shared), zipper_replace(original_pattern.to_expr(&shared), &shared.zippers.zip_of_zid[hole_zid], Node::Prim(format!("<{expands_to}>").into()))) } continue 'expansion; } // Pruning (FREE VARS): if an invention has free variables in the body then it's not a real function and we can discard it // Here we just check if our expansion just yielded a variable, and if that is bound based on how many lambdas there are above it. - if expands_to.free_variable(shared.zip_of_zid[hole_zid].depth_root_to_arg()) { + if expands_to.free_variable(shared.zippers.zip_of_zid[hole_zid].depth_root_to_arg()) { if !shared.cfg.no_stats { shared.stats.lock().deref_mut().free_vars_fired += 1; }; if tracked && !shared.cfg.quiet { println!("{} pruned by free var in body when expanding {} to {}", "[TRACK]".red().bold(), original_pattern.to_expr(&shared), original_pattern.show_track_expansion(hole_zid, &shared)) } continue 'expansion; // free var @@ -1048,7 +1049,7 @@ fn stitch_search( } if !shared.cfg.no_stats { shared.stats.lock().calc_unargcap += 1; }; - inverse_argument_capture(&mut finished_pattern, &shared.cfg, &shared.zip_of_zid, &shared.arg_of_zid_node, &shared.extensions_of_zid, &shared.set, &shared.analyzed_ivars, &shared.cost_fn); + inverse_argument_capture(&mut finished_pattern, &shared.cfg, &shared.zippers, &shared.extensions_of_zid, &shared.set, &shared.analyzed_ivars, &shared.cost_fn); // Pruning (UPPER BOUND) if finished_pattern.utility <= weak_utility_pruning_cutoff { @@ -1084,7 +1085,7 @@ fn stitch_search( if original_pattern.tracked && !found_tracked { // let new = format!("<{}>",tracked_expands_to(&original_pattern, hole_zid, &shared)); - // let mut s = original_pattern.to_expr(&shared).zipper_replace(&shared.zip_of_zid[hole_zid], &new ).to_string(); + // let mut s = original_pattern.to_expr(&shared).zipper_replace(&shared.zippers.zip_of_zid[hole_zid], &new ).to_string(); // s = s.replace(&new, &new.clone().magenta().bold().to_string()); if !shared.cfg.quiet { println!("{} pruned when expanding because there were no match locations for the target expansion of {} to {}", "[TRACK]".red().bold(), original_pattern.to_expr(&shared), original_pattern.show_track_expansion(hole_zid, &shared)) } } @@ -1185,16 +1186,12 @@ fn get_zippers( analyzed_cost: &AnalyzedExpr, set: &mut ExprSet, analyzed_free_vars: &mut AnalyzedExpr, -) -> (FxHashMap, Vec, Vec>, FxHashMap>, Vec) { +) -> (Zippers, FxHashMap>, Vec) { - let mut zid_of_zip: FxHashMap = Default::default(); - let mut zip_of_zid: Vec = Default::default(); - let mut arg_of_zid_node: Vec> = Default::default(); + let mut zippers = Zippers::default(); let mut zids_of_node: FxHashMap> = Default::default(); - zid_of_zip.insert(Zipper::default(), EMPTY_ZID); - zip_of_zid.push(Zipper::default()); - arg_of_zid_node.push(FxHashMap::default()); + zippers.add_empty(EMPTY_ZID); // loop over all nodes in all programs in bottom up order for idx in corpus_span.clone() { @@ -1207,8 +1204,7 @@ fn get_zippers( // clone to appease the borrow checker let node = set.get(idx).node().clone(); - arg_of_zid_node[EMPTY_ZID].insert(idx, - Arg { shifted_id: idx, unshifted_id: idx, shift: 0, cost: analyzed_cost[idx] as Cost, expands_to: expands_to_of_node(&node) }); + zippers.add_arg(EMPTY_ZID, idx, analyzed_cost[idx] as Cost, expands_to_of_node(&node)); match node { Node::IVar(_) => { unreachable!() } @@ -1217,100 +1213,30 @@ fn get_zippers( // bubble from `f` for f_zid in zids_of_node[&f].iter() { // clone and extend zip to get new zid for this node - let mut zip = zip_of_zid[*f_zid].clone(); - zip.add_to_front(ZNode::Func); - let zid = zid_of_zip.entry(zip.clone()).or_insert_with(|| { - let zid = zip_of_zid.len(); - zip_of_zid.push(zip); - arg_of_zid_node.push(FxHashMap::default()); - zid - }); - // add new zid to this node - zids.push(*zid); - // give it the same arg - let arg = arg_of_zid_node[*f_zid][&f].clone(); - arg_of_zid_node[*zid].insert(idx, arg); + zids.push(zippers.extend_zipper(*f_zid, idx, f, ZNode::Func)) } // bubble from `x` for x_zid in zids_of_node[&x].iter() { - // clone and extend zip to get new zid for this node - let mut zip = zip_of_zid[*x_zid].clone(); - zip.add_to_front(ZNode::Arg); - let zid = zid_of_zip.entry(zip.clone()).or_insert_with(|| { - let zid = zip_of_zid.len(); - zip_of_zid.push(zip); - arg_of_zid_node.push(FxHashMap::default()); - zid - }); - // add new zid to this node - zids.push(*zid); - // give it the same arg - let arg = arg_of_zid_node[*x_zid][&x].clone(); - arg_of_zid_node[*zid].insert(idx, arg); - + zids.push(zippers.extend_zipper(*x_zid, idx, x, ZNode::Arg)) } }, Node::Lam(b, _) => { for b_zid in zids_of_node[&b].iter() { - - // clone and extend zip to get new zid for this node - let mut zip = zip_of_zid[*b_zid].clone(); - zip.add_to_front(ZNode::Body); - let zid = zid_of_zip.entry(zip.clone()).or_insert_with(|| { - let zid = zip_of_zid.len(); - zip_of_zid.push(zip.clone()); - arg_of_zid_node.push(FxHashMap::default()); - zid - }); - // add new zid to this node - zids.push(*zid); - // shift the arg but keep the unshifted part the same - let mut arg: Arg = arg_of_zid_node[*b_zid][&b].clone(); - - if !analyzed_free_vars.analyze_get(set.get(arg.shifted_id)).is_empty() { - // the arg has free vars so we should actually downshift it by 1 - if analyzed_free_vars[arg.shifted_id].contains(&0) { - // furthermore one of those vars is a 0 then it will get shifted to -1, so we handle that slightly specially - // by inserting an IVar to indicate this - - // how many lambdas are along this zipper? (including most recent one) - let depth_root_to_arg = zip.depth_root_to_arg() as i32; - - // find all pointers to $0 (this is the `init_depth` parameter) and replace then with #(num_lams - 1) that is - // point past all lambdas except the newly added one. For example if there were no lambdas other than the - // newly added one this would be num_lams=1 so it'd be #0. - arg.shifted_id = insert_arg_ivars(&mut set.get_mut(arg.shifted_id), depth_root_to_arg-1, 0, analyzed_free_vars); - } - arg.shifted_id = set.get_mut(arg.shifted_id).shift(-1, 0, analyzed_free_vars); - arg.shift -= 1; - } - arg_of_zid_node[*zid].insert(idx, arg); + let zid = zippers.extend_zipper(*b_zid, idx, b, ZNode::Body); + zids.push(zid); + zippers.handle_shift(zid, *b_zid, idx, b, analyzed_free_vars, set); + } }, } zids_of_node.insert(idx, zids); } - let extensions_of_zid = zip_of_zid.iter().map(|zip| { - let mut zip_body = zip.clone(); - zip_body.add_to_end(ZNode::Body); - let mut zip_arg = zip.clone(); - zip_arg.add_to_end(ZNode::Arg); - let mut zip_func = zip.clone(); - zip_func.add_to_end(ZNode::Func); - ZIdExtension { - body: zid_of_zip.get(&zip_body).copied(), - arg: zid_of_zip.get(&zip_arg).copied(), - func: zid_of_zip.get(&zip_func).copied(), - } - }).collect(); + let extensions_of_zid = zippers.compute_extensions(); + - (zid_of_zip, - zip_of_zid, - arg_of_zid_node, - zids_of_node, - extensions_of_zid) + (zippers, zids_of_node, extensions_of_zid) } /// the complete result of a single step of compression, this is a somewhat expensive data structure @@ -1549,7 +1475,7 @@ fn get_utility_of_loc_once(pattern: &Pattern, shared: &SharedData) -> Vec // extra utility. Note that it doesn't matter which // of the zids we use as long as it corresponds to the right ivar let multiuse_utility = ivar_multiuses.iter().map(|(zid,count)| - count * shared.arg_of_zid_node[*zid][loc].cost as Cost + count * shared.zippers.arg_of_zid_node[*zid][loc].cost as Cost ).sum::(); // if !shared.cfg.quiet { println!("multiuse {}", multiuse_utility) } @@ -1579,7 +1505,7 @@ fn bottom_up_utility_correction(pattern: &Pattern, shared:&SharedData, utility_o // this node is a potential rewrite location let utility_of_args: Cost = pattern.pattern_args.iterate_one_zid_per_argument() - .map(|zid| cumulative_utility_of_node[shared.arg_of_zid_node[zid][&node].unshifted_id]) + .map(|zid| cumulative_utility_of_node[shared.zippers.arg_of_zid_node[zid][&node].unshifted_id]) .sum(); let utility_with_rewrite = utility_of_args + utility_of_loc_once[idx]; @@ -1613,7 +1539,7 @@ pub fn inverse_delta(cost_once: Cost, usages: Cost, arg_uses: usize, cost_fn: &E // (not used in popl code - experimental; always exists at the first return statement unless --inv-arg-cap is turned on) #[allow(clippy::too_many_arguments)] -pub fn inverse_argument_capture(finished: &mut FinishedPattern, cfg: &CompressionStepConfig, zip_of_zid: &[Zipper], arg_of_zid_node: &[FxHashMap], extensions_of_zid: &[ZIdExtension], set: &ExprSet, analyzed_ivars: &AnalyzedExpr, cost_fn: &ExprCost) { +pub fn inverse_argument_capture(finished: &mut FinishedPattern, cfg: &CompressionStepConfig, zippers: &Zippers, extensions_of_zid: &[ZIdExtension], set: &ExprSet, analyzed_ivars: &AnalyzedExpr, cost_fn: &ExprCost) { if !cfg.inv_arg_cap || cfg.no_other_util { return } @@ -1624,7 +1550,7 @@ pub fn inverse_argument_capture(finished: &mut FinishedPattern, cfg: &Compressio let _max_num_to_add = cfg.max_arity - finished.arity; while finished.arity < cfg.max_arity { - let counts = use_counts(&finished.pattern, zip_of_zid, arg_of_zid_node, extensions_of_zid, set, analyzed_ivars); + let counts = use_counts(&finished.pattern, zippers, extensions_of_zid, set, analyzed_ivars); let possible_to_uninline = possible_to_uninline(counts, finished.usages, cost_fn); let best = possible_to_uninline.into_iter().max_by_key(|(delta, _compressive_delta, _noncompressive_delta, _cost, _zids)| *delta); @@ -1665,14 +1591,14 @@ fn possible_to_uninline(counts: FxHashMap)>, finished_usa } /// not used in popl code - experimental -fn use_counts(pattern: &Pattern, zip_of_zid: &[Zipper], arg_of_zid_node: &[FxHashMap], extensions_of_zid: &[ZIdExtension], set: &ExprSet, analyzed_ivars: &AnalyzedExpr) -> FxHashMap)> { +fn use_counts(pattern: &Pattern, zippers: &Zippers, extensions_of_zid: &[ZIdExtension], set: &ExprSet, analyzed_ivars: &AnalyzedExpr) -> FxHashMap)> { let mut curr_zip: Zipper = Zipper::default(); let curr_zid: ZId = EMPTY_ZID; let zids = &pattern.pattern_args.iterate_arguments().cloned().collect::>(); // map zids to zips with a bool thats true if this is a hole and false if its a future ivar let zips: Vec = zids.iter() - .map(|labelled_zid| zip_of_zid[labelled_zid.zid].clone()).collect(); + .map(|labelled_zid| zippers.zip_of_zid[labelled_zid.zid].clone()).collect(); let mut counts: FxHashMap)> = Default::default(); @@ -1714,7 +1640,7 @@ fn use_counts(pattern: &Pattern, zip_of_zid: &[Zipper], arg_of_zid_node: &[FxHas } } // we can pick any match location - helper(pattern.match_locations[0], pattern.match_locations[0], &mut curr_zip, curr_zid, &zips, zids, arg_of_zid_node, extensions_of_zid, set, &mut counts, analyzed_ivars); + helper(pattern.match_locations[0], pattern.match_locations[0], &mut curr_zip, curr_zid, &zips, zids, &zippers.arg_of_zid_node, extensions_of_zid, set, &mut counts, analyzed_ivars); counts } @@ -1890,17 +1816,14 @@ pub fn construct_shared( if !cfg.quiet { println!("cost_of_node structs: {:?}ms", tstart.elapsed().as_millis()) } tstart = std::time::Instant::now(); - let (zid_of_zip, - zip_of_zid, - arg_of_zid_node, + let (zippers, zids_of_node, extensions_of_zid) = get_zippers(&corpus_span, &analyzed_cost, &mut set, &mut analyzed_free_vars); if !cfg.quiet { println!("get_zippers(): {:?}ms", tstart.elapsed().as_millis()) } tstart = std::time::Instant::now(); - if !cfg.quiet { println!("{} zips", zip_of_zid.len()) } - if !cfg.quiet { println!("arg_of_zid_node size: {}", arg_of_zid_node.len()) } + if !cfg.quiet { zippers.print_stats() } // set up tracking if any let tracking: Option = { @@ -1908,7 +1831,7 @@ pub fn construct_shared( let mut set = ExprSet::empty(Order::ChildFirst, false, false); let idx = set.parse_extend(s).unwrap(); let expr = ExprOwned::new(set,idx); - if let Some(zids_of_ivar) = zids_of_ivar_of_expr(&expr, &zid_of_zip) { + if let Some(zids_of_ivar) = zids_of_ivar_of_expr(&expr, &zippers) { Some(Tracking { expr, zids_of_ivar }) } else { if !cfg.quiet { println!("Tracking: can't possibly find a match for this in corpus because one if the necessary zippers ZIDs doesnt exist in corpus")} @@ -2018,7 +1941,7 @@ pub fn construct_shared( }; // This handle the case covered by Appendix B in the paper - inverse_argument_capture(&mut finished_pattern, cfg, &zip_of_zid, &arg_of_zid_node, &extensions_of_zid, &set, &analyzed_ivars, cost_fn); + inverse_argument_capture(&mut finished_pattern, cfg, &zippers, &extensions_of_zid, &set, &analyzed_ivars, cost_fn); if !cfg.no_stats { stats.azero_calc_unargcap += 1; }; // Pruning (UPPER BOUND): This is the full upper bound pruning @@ -2056,7 +1979,7 @@ pub fn construct_shared( let shared = Arc::new(SharedData { crit: Mutex::new(crit), programs: programs.to_vec(), - arg_of_zid_node, + zippers, cost_fn: cost_fn.clone(), analyzed_free_vars, analyzed_ivars, @@ -2065,8 +1988,6 @@ pub fn construct_shared( corpus_span: corpus_span.clone(), roots: roots.to_vec(), zids_of_node, - zip_of_zid, - zid_of_zip, extensions_of_zid, set, num_paths_to_node, diff --git a/src/expansion.rs b/src/expansion.rs index 3c853f3d..6a825d08 100644 --- a/src/expansion.rs +++ b/src/expansion.rs @@ -134,7 +134,7 @@ impl std::fmt::Display for ExpandsTo { pub fn tracked_expands_to(pattern: &Pattern, hole_zid: ZId, shared: &SharedData) -> ExpandsTo { // apply the hole zipper to the original expr being tracked to get the subtree // this will expand into, then get the ExpandsTo of that - let idx = shared.zip_of_zid[hole_zid].zip(&shared.tracking.as_ref().unwrap().expr); + let idx = shared.zippers.zip_of_zid[hole_zid].zip(&shared.tracking.as_ref().unwrap().expr); match expands_to_of_node(&shared.tracking.as_ref().unwrap().expr.set[idx]) { ExpandsTo(ExpandsToInner::IVar(i, VariableType::Metavar)) => { ExpandsTo(ExpandsToInner::IVar(pattern.pattern_args.find_variable(shared, i as usize) as i32, VariableType::Metavar)) @@ -176,7 +176,7 @@ pub fn get_ivars_expansions(original_pattern: &Pattern, arg_of_loc: &FxHashMap Vec { self.variables.iter().map(|(zid, _)| - shared.arg_of_zid_node[*zid as usize][node].shifted_id + shared.zippers.arg_of_zid_node[*zid as usize][node].shifted_id ).collect() } @@ -77,7 +77,7 @@ impl PatternArgs { continue; } let zid = self.variables[i].0 as ZId; - let shifted_arg = shared.arg_of_zid_node[zid][loc].shifted_id; + let shifted_arg = shared.zippers.arg_of_zid_node[zid][loc].shifted_id; if !shared.analyzed_ivars[shifted_arg].is_empty() { return true; } @@ -95,8 +95,8 @@ impl PatternArgs { continue; } // if its the same arg in every place, and doesnt have any free vars (ie it's safe to inline) - if locs.iter().map(|loc| shared.arg_of_zid_node[argchoice.zid][loc].shifted_id).all_equal() - && shared.analyzed_free_vars[shared.arg_of_zid_node[argchoice.zid][&locs[0]].shifted_id].is_empty() + if locs.iter().map(|loc| shared.zippers.arg_of_zid_node[argchoice.zid][loc].shifted_id).all_equal() + && shared.analyzed_free_vars[shared.zippers.arg_of_zid_node[argchoice.zid][&locs[0]].shifted_id].is_empty() { if !shared.cfg.no_stats { shared.stats.lock().deref_mut().useless_abstract_fired += 1; }; return true; @@ -117,13 +117,13 @@ impl PatternArgs { if *type_1 != VariableType::Metavar { continue; } - let arg_of_loc_1 = &shared.arg_of_zid_node[*ivar_zid_1 as ZId]; + let arg_of_loc_1 = &shared.zippers.arg_of_zid_node[*ivar_zid_1 as ZId]; // for some reason, the enumerate makes it like 1% faster????? for (_j, (ivar_zid_2, type_2)) in self.variables.iter().enumerate().skip(i+1) { if *type_2 != VariableType::Metavar { continue; } - let arg_of_loc_2 = &shared.arg_of_zid_node[*ivar_zid_2 as ZId]; + let arg_of_loc_2 = &shared.zippers.arg_of_zid_node[*ivar_zid_2 as ZId]; if locs.iter().all(|loc| arg_of_loc_1[loc].shifted_id == arg_of_loc_2[loc].shifted_id) { @@ -192,7 +192,7 @@ impl LocationsForReusableArgs<'_> { impl PatternArgs { pub fn reusable_args_location(&self, shared: &SharedData, ivar: Idx, arg_of_loc: &FxHashMap, match_locations: &mut LocationsForReusableArgs) -> Vec { let (first_zid_of_var, type_of_var) = self.variables[ivar]; - let arg_of_loc_ivar = &shared.arg_of_zid_node[first_zid_of_var as ZId]; + let arg_of_loc_ivar = &shared.zippers.arg_of_zid_node[first_zid_of_var as ZId]; let relevant_locs = match_locations.relevant_locs(type_of_var, arg_of_loc, &shared.sym_var_info); compatible_locations(shared, relevant_locs, arg_of_loc, arg_of_loc_ivar, type_of_var) } diff --git a/src/rewriting.rs b/src/rewriting.rs index f88dd2c7..1ee7a657 100644 --- a/src/rewriting.rs +++ b/src/rewriting.rs @@ -39,13 +39,13 @@ pub fn rewrite_fast( && (!pattern.util_calc.corrected_utils.contains_key(&unshifted_id) // and either we have no conflict (ie corrected_utils doesnt have an entry) || pattern.util_calc.corrected_utils[&unshifted_id]) // or we have a conflict but we choose to accept it (which is contextless in this top down approach so its the right move) // && !pattern.pattern.variables.iter().any(|zid| // and there are no negative vars anywhere in the arguments - // shared.egraph[shared.arg_of_zid_node[*zid][&unshifted_id].Idx].data.free_vars.iter().any(|var| *var < 0)) + // shared.egraph[shared.zippers.arg_of_zid_node[*zid][&unshifted_id].Idx].data.free_vars.iter().any(|var| *var < 0)) { // if !shared.cfg.quiet { println!("inv applies at unshifted={} with shift={}", extract(unshifted_id,&shared.egraph), shift) } let mut expr = owned_set.add(inv_name.clone()); // wrap the prim in all the Apps to args for zid in pattern.pattern.pattern_args.iterate_one_zid_per_argument() { - let arg: &Arg = &shared.arg_of_zid_node[zid][&unshifted_id]; + let arg: &Arg = &shared.zippers.arg_of_zid_node[zid][&unshifted_id]; if arg.shift != 0 { shift_rules.push(ShiftRule{depth_cutoff: total_depth, shift: arg.shift}); @@ -66,7 +66,7 @@ pub fn rewrite_fast( // Also note that in the single_hole code --eta-long enforces that match locations never contains anything that starts to the left of a func so // we dont need to worry about the case where the zipper would extend even past the root of the match location // Also note that due to beta normal form, this will be zero and will be a no-op if the arg is a lambda - let arity_of_arg = shared.zip_of_zid[zid].function_arity(); + let arity_of_arg = shared.zippers.zip_of_zid[zid].function_arity(); if arity_of_arg > 0 { let analyzed_free_vars = &mut AnalyzedExpr::new(FreeVarAnalysis); diff --git a/src/tdfa.rs b/src/tdfa.rs index 0cf22eb2..a806022b 100644 --- a/src/tdfa.rs +++ b/src/tdfa.rs @@ -173,7 +173,7 @@ impl TDFAInventionAnnotation { let root_sym = global_annotations.symbols[match_location].clone()?; let mut ivar_states: Vec = vec![]; let all_found = pattern.pattern_args.iterate_one_zid_per_argument().all(|ivar_zid| { - let Some(node) = shared.arg_of_zid_node[ivar_zid].get(&match_location) else { + let Some(node) = shared.zippers.arg_of_zid_node[ivar_zid].get(&match_location) else { return false; }; let Some(ivar_sym) = global_annotations.symbols[node.unshifted_id].clone() else { diff --git a/src/zipper.rs b/src/zipper.rs index 1dc5d0e4..18f775b1 100644 --- a/src/zipper.rs +++ b/src/zipper.rs @@ -1,5 +1,8 @@ -use lambdas::ExprOwned; +use lambdas::{AnalyzedExpr, ExprOwned, ExprSet, FreeVarAnalysis, Idx}; pub use lambdas::{ZNode, ZId, LabelledZId}; +use rustc_hash::FxHashMap; + +use crate::{insert_arg_ivars, Arg, Cost, ExpandsTo, ZIdExtension}; #[derive(Clone, Debug, PartialEq, Eq, Hash, Default)] pub struct Zipper(Vec); @@ -47,4 +50,100 @@ impl Zipper { expr.immut().zip(&self.0).idx } +} + +#[derive(Clone, Debug, Default)] +pub struct Zippers { + pub zid_of_zip: FxHashMap, + pub zip_of_zid: Vec, + pub arg_of_zid_node: Vec>, +} + +impl Zippers { + + #[inline(always)] + pub fn get_interned_idx(&self, zipper: &Zipper) -> Option { + self.zid_of_zip.get(zipper).cloned() + } + + #[inline(always)] + pub fn add_empty(&mut self, empty_zid: ZId) { + self.zid_of_zip.insert(Zipper::default(), empty_zid); + self.zip_of_zid.push(Zipper::default()); + self.arg_of_zid_node.push(FxHashMap::default()); + } + + #[inline(always)] + pub fn add_arg(&mut self, zid: ZId, node: Idx, cost: Cost, expands_to: ExpandsTo) { + self.arg_of_zid_node[zid].insert(node, + Arg { shifted_id: node, unshifted_id: node, shift: 0, cost, expands_to }); + + } + + #[inline(always)] + pub fn extend_zipper(&mut self, unextended_zid: ZId, extended_node: Idx, unextended_node: Idx, znode: ZNode) -> usize { + let mut zip = self.zip_of_zid[unextended_zid].clone(); + zip.add_to_front(znode); + let zip_of_zid = &mut self.zip_of_zid; + let arg_of_zid_node = &mut self.arg_of_zid_node; + let zid = self.zid_of_zip.entry(zip.clone()).or_insert_with(|| { + let zid = zip_of_zid.len(); + zip_of_zid.push(zip); + arg_of_zid_node.push(FxHashMap::default()); + zid + }); + // add new zid to this node + // give it the same arg + let arg = self.arg_of_zid_node[unextended_zid][&unextended_node].clone(); + self.arg_of_zid_node[*zid].insert(extended_node, arg); + *zid + } + + #[inline(always)] + pub fn handle_shift(&mut self, extended_zid: ZId, unextended_zid: ZId, extended_node: Idx, unextended_node: Idx, analyzed_free_vars: &mut AnalyzedExpr, set: &mut ExprSet) { + let zip = &self.zip_of_zid[extended_zid]; + let mut arg: Arg = self.arg_of_zid_node[unextended_zid][&unextended_node].clone(); + // shift the arg but keep the unshifted part the same + if !analyzed_free_vars.analyze_get(set.get(arg.shifted_id)).is_empty() { + // the arg has free vars so we should actually downshift it by 1 + if analyzed_free_vars[arg.shifted_id].contains(&0) { + // furthermore one of those vars is a 0 then it will get shifted to -1, so we handle that slightly specially + // by inserting an IVar to indicate this + + // how many lambdas are along this zipper? (including most recent one) + let depth_root_to_arg = zip.depth_root_to_arg() as i32; + + // find all pointers to $0 (this is the `init_depth` parameter) and replace then with #(num_lams - 1) that is + // point past all lambdas except the newly added one. For example if there were no lambdas other than the + // newly added one this would be num_lams=1 so it'd be #0. + arg.shifted_id = insert_arg_ivars(&mut set.get_mut(arg.shifted_id), depth_root_to_arg-1, 0, analyzed_free_vars); + } + arg.shifted_id = set.get_mut(arg.shifted_id).shift(-1, 0, analyzed_free_vars); + arg.shift -= 1; + } + self.arg_of_zid_node[extended_zid].insert(extended_node, arg); + } + + #[inline(always)] + pub fn compute_extensions(&self) -> Vec { + self.zip_of_zid.iter().map(|zip| { + let mut zip_body = zip.clone(); + zip_body.add_to_end(ZNode::Body); + let mut zip_arg = zip.clone(); + zip_arg.add_to_end(ZNode::Arg); + let mut zip_func = zip.clone(); + zip_func.add_to_end(ZNode::Func); + ZIdExtension { + body: self.zid_of_zip.get(&zip_body).copied(), + arg: self.zid_of_zip.get(&zip_arg).copied(), + func: self.zid_of_zip.get(&zip_func).copied(), + } + }).collect() + } + + pub fn print_stats(&self) { + println!("{} zips", self.zip_of_zid.len()); + println!("arg_of_zid_node size: {}", self.arg_of_zid_node.len()) + } + } \ No newline at end of file