Add action to set custom cost

src/actions.rs

@@ -39,6 +39,18 @@ impl<'a> ActionCompiler<'a> {
                 self.do_atom_term(e);
                 self.instructions.push(Instruction::Set(func.name));
             }
+            GenericCoreAction::Cost(_ann, f, args, e) => {
+                let ResolvedCall::Func(func) = f else {
+                    panic!(
+                        "Cannot set cost of primitive- should have been caught by typechecking!!!"
+                    )
+                };
+                for arg in args {
+                    self.do_atom_term(arg);
+                }
+                self.do_atom_term(e);
+                self.instructions.push(Instruction::Cost(func.name));
+            }
             GenericCoreAction::Change(_ann, change, f, args) => {
                 let ResolvedCall::Func(func) = f else {
                     panic!("Cannot change primitive- should have been caught by typechecking!!!")
@@ -128,12 +140,15 @@ enum Instruction {
     /// Pop primitive arguments off the stack, calls the primitive,
     /// and push the result onto the stack.
     CallPrimitive(SpecializedPrimitive, usize),
-    /// Pop function arguments off the stack and either deletes or subsumes the corresponding row
-    /// in the function.
+    /// Pop function arguments off the stack and either deletes, subsumes, or changes the cost
+    /// of the corresponding row in the function.
     Change(Change, Symbol),
     /// Pop the value to be set and the function arguments off the stack.
     /// Set the function at the given arguments to the new value.
     Set(Symbol),
+    /// Pop the value to have its cost set and the function arguments off the stack.
+    /// Set the function at the given arguments to the new cost.
+    Cost(Symbol),
     /// Union the last `n` values on the stack.
     Union(usize, ArcSort),
     /// Extract the best expression. `n` is always 2.
@@ -346,10 +361,25 @@ impl EGraph {
                     // set to union
                     let new_value = stack.pop().unwrap();
                     let new_len = stack.len() - function.schema.input.len();
-
                     self.perform_set(*f, new_value, stack)?;
                     stack.truncate(new_len)
                 }
+                Instruction::Cost(f) => {
+                    let function = self.functions.get_mut(f).unwrap();
+                    let new_cost = stack.pop().unwrap();
+                    let new_len = stack.len() - function.schema.input.len();
+
+                    let function = self.functions.get_mut(f).unwrap();
+
+                    let args = &stack[new_len..];
+
+                    let i64sort: Arc<I64Sort> = self.type_info.get_sort_nofail();
+                    let cost = i64::load(&i64sort, &new_cost);
+                    function.update_cost(args, cost.try_into().unwrap());
+
+                    stack.truncate(new_len);
+                }
+
                 Instruction::Union(arity, sort) => {
                     let new_len = stack.len() - arity;
                     let values = &stack[new_len..];

src/ast/mod.rs

@@ -1222,6 +1222,13 @@ where
         Vec<GenericExpr<Head, Leaf>>,
         GenericExpr<Head, Leaf>,
     ),
+    /// `cost` sets the cost of a function to a particular value.
+    Cost(
+        Span,
+        Head,
+        Vec<GenericExpr<Head, Leaf>>,
+        GenericExpr<Head, Leaf>,
+    ),
     /// Delete or subsume (mark as hidden from future rewrites and unextractable) an entry from a function.
     Change(Span, Change, Head, Vec<GenericExpr<Head, Leaf>>),
     /// `union` two datatypes, making them equal
@@ -1304,6 +1311,7 @@ where
         match self {
             GenericAction::Let(_ann, lhs, rhs) => list!("let", lhs, rhs),
             GenericAction::Set(_ann, lhs, args, rhs) => list!("set", list!(lhs, ++ args), rhs),
+            GenericAction::Cost(_ann, lhs, args, rhs) => list!("cost", list!(lhs, ++ args), rhs),
             GenericAction::Union(_ann, lhs, rhs) => list!("union", lhs, rhs),
             GenericAction::Change(_ann, change, lhs, args) => {
                 list!(
@@ -1344,6 +1352,15 @@ where
                     right,
                 )
             }
+            GenericAction::Cost(span, lhs, args, rhs) => {
+                let right = f(rhs);
+                GenericAction::Cost(
+                    span.clone(),
+                    lhs.clone(),
+                    args.iter().map(f).collect(),
+                    right,
+                )
+            }
             GenericAction::Change(span, change, lhs, args) => GenericAction::Change(
                 span.clone(),
                 *change,
@@ -1378,6 +1395,10 @@ where
                 let args = args.into_iter().map(|e| e.visit_exprs(f)).collect();
                 GenericAction::Set(span, lhs.clone(), args, rhs.visit_exprs(f))
             }
+            GenericAction::Cost(span, lhs, args, rhs) => {
+                let args = args.into_iter().map(|e| e.visit_exprs(f)).collect();
+                GenericAction::Cost(span, lhs.clone(), args, rhs.visit_exprs(f))
+            }
             GenericAction::Change(span, change, lhs, args) => {
                 let args = args.into_iter().map(|e| e.visit_exprs(f)).collect();
                 GenericAction::Change(span, change, lhs.clone(), args)
@@ -1417,6 +1438,14 @@ where
                 let rhs = rhs.subst_leaf(&mut fvar_expr!());
                 GenericAction::Set(span, lhs.clone(), args, rhs)
             }
+            GenericAction::Cost(span, lhs, args, rhs) => {
+                let args = args
+                    .into_iter()
+                    .map(|e| e.subst_leaf(&mut fvar_expr!()))
+                    .collect();
+                let rhs = rhs.subst_leaf(&mut fvar_expr!());
+                GenericAction::Cost(span, lhs.clone(), args, rhs)
+            }
             GenericAction::Change(span, change, lhs, args) => {
                 let args = args
                     .into_iter()

src/ast/parse.rs

@@ -621,6 +621,12 @@ fn non_let_action(ctx: &Context) -> Res<Action> {
             expr,
         ))
         .map(|((), (f, args), v), span| Action::Set(span, f, args, v))(ctx),
+        "cost" => parens(sequence3(
+            text("cost"),
+            parens(sequence(ident, repeat_until_end_paren(expr))),
+            expr,
+        ))
+        .map(|((), (f, args), v), span| Action::Cost(span, f, args, v))(ctx),
         "delete" => parens(sequence(
             text("delete"),
             parens(sequence(ident, repeat_until_end_paren(expr))),

src/constraint.rs

@@ -365,6 +365,32 @@ impl Assignment<AtomTerm, ArcSort> {
                     rhs,
                 ))
             }
+            // (cost (f [*x]) rhs) where rhs should have type I64
+            GenericAction::Cost(
+                span,
+                CorrespondingVar {
+                    head,
+                    to: _mapped_var,
+                },
+                children,
+                rhs,
+            ) => {
+                let children: Vec<_> = children
+                    .iter()
+                    .map(|child| self.annotate_expr(child, typeinfo))
+                    .collect();
+                let types: Vec<_> = children.iter().map(|child| child.output_type()).collect();
+                let resolved_call =
+                    ResolvedCall::from_resolution_func_types(head, &types, typeinfo)
+                        .ok_or_else(|| TypeError::UnboundFunction(*head, span.clone()))?;
+                let resolved_rhs = self.annotate_expr(rhs, typeinfo);
+                Ok(ResolvedAction::Cost(
+                    span.clone(),
+                    resolved_call,
+                    children.clone(),
+                    resolved_rhs,
+                ))
+            }
             // Note mapped_var for delete is a dummy variable that does not mean anything
             GenericAction::Change(
                 span,
@@ -541,6 +567,24 @@ impl CoreAction {
                     )?)
                     .collect())
             }
+            CoreAction::Cost(span, head, args, rhs) => {
+                let mut args = args.clone();
+                let var = symbol_gen.fresh(head);
+                args.push(AtomTerm::Var(span.clone(), var));
+
+                let mut all_terms = args.clone();
+                all_terms.push(rhs.clone());
+
+                Ok(get_literal_and_global_constraints(&all_terms, typeinfo)
+                    .chain(get_atom_application_constraints(
+                        head, &args, span, typeinfo,
+                    )?)
+                    .chain(once(Constraint::Assign(
+                        rhs.clone(),
+                        typeinfo.get_sort_nofail::<I64Sort>() as ArcSort,
+                    )))
+                    .collect())
+            }
             CoreAction::Change(span, _change, head, args) => {
                 let mut args = args.clone();
                 // Add a dummy last output argument

src/core.rs

@@ -384,6 +384,12 @@ pub enum GenericCoreAction<Head, Leaf> {
         Vec<GenericAtomTerm<Leaf>>,
         GenericAtomTerm<Leaf>,
     ),
+    Cost(
+        Span,
+        Head,
+        Vec<GenericAtomTerm<Leaf>>,
+        GenericAtomTerm<Leaf>,
+    ),
     Change(Span, Change, Head, Vec<GenericAtomTerm<Leaf>>),
     Union(Span, GenericAtomTerm<Leaf>, GenericAtomTerm<Leaf>),
     Panic(Span, String),
@@ -492,6 +498,34 @@ where
                         mapped_expr,
                     ));
                 }
+                GenericAction::Cost(span, head, args, expr) => {
+                    let mut mapped_args = vec![];
+                    for arg in args {
+                        let (actions, mapped_arg) =
+                            arg.to_core_actions(typeinfo, binding, fresh_gen)?;
+                        norm_actions.extend(actions.0);
+                        mapped_args.push(mapped_arg);
+                    }
+                    let (actions, mapped_expr) =
+                        expr.to_core_actions(typeinfo, binding, fresh_gen)?;
+                    norm_actions.extend(actions.0);
+                    norm_actions.push(GenericCoreAction::Cost(
+                        span.clone(),
+                        head.clone(),
+                        mapped_args
+                            .iter()
+                            .map(|e| e.get_corresponding_var_or_lit(typeinfo))
+                            .collect(),
+                        mapped_expr.get_corresponding_var_or_lit(typeinfo),
+                    ));
+                    let v = fresh_gen.fresh(head);
+                    mapped_actions.0.push(GenericAction::Cost(
+                        span.clone(),
+                        CorrespondingVar::new(head.clone(), v),
+                        mapped_args,
+                        mapped_expr,
+                    ));
+                }
                 GenericAction::Change(span, change, head, args) => {
                     let mut mapped_args = vec![];
                     for arg in args {

src/extract.rs

@@ -166,8 +166,9 @@ impl<'a> Extractor<'a> {
         function: &Function,
         children: &[Value],
         termdag: &mut TermDag,
+        cost: Option<usize>,
     ) -> Option<(Vec<Term>, Cost)> {
-        let mut cost = function.decl.cost.unwrap_or(1);
+        let mut cost = cost.unwrap_or(function.decl.cost.unwrap_or(1));
         let types = &function.schema.input;
         let mut terms: Vec<Term> = vec![];
         for (ty, value) in types.iter().zip(children) {
@@ -188,7 +189,7 @@ impl<'a> Extractor<'a> {
                 if func.schema.output.is_eq_sort() {
                     for (inputs, output) in func.nodes.iter(false) {
                         if let Some((term_inputs, new_cost)) =
-                            self.node_total_cost(func, inputs, termdag)
+                            self.node_total_cost(func, inputs, termdag, output.cost)
                         {
                             let make_new_pair = || (new_cost, termdag.app(sym, term_inputs));
 

src/function/mod.rs

@@ -45,6 +45,7 @@ pub struct TupleOutput {
     pub value: Value,
     pub timestamp: u32,
     pub subsumed: bool,
+    pub cost: Option<usize>,
 }
 
 #[derive(Clone, Debug)]
@@ -240,6 +241,11 @@ impl Function {
         self.nodes.get_mut(inputs).unwrap().subsumed = true;
     }
 
+    /// Updates the cost of the given inputs.
+    pub fn update_cost(&mut self, inputs: &[Value], cost: usize) {
+        self.nodes.get_mut(inputs).unwrap().cost = Some(cost);
+    }
+
     /// Return a column index that contains (a superset of) the offsets for the
     /// given column. This method can return nothing if the indexes available
     /// contain too many irrelevant offsets.
@@ -443,7 +449,7 @@ impl Function {
         }
         let out_ty = &self.schema.output;
         self.nodes
-            .insert_and_merge(scratch, timestamp, out.subsumed, |prev| {
+            .insert_and_merge(scratch, timestamp, out.subsumed, out.cost, |prev| {
                 if let Some(mut prev) = prev {
                     out_ty.canonicalize(&mut prev, uf);
                     let mut appended = false;

src/function/table.rs

@@ -135,7 +135,7 @@ impl Table {
     /// previous value, if there was one.
     pub(crate) fn insert(&mut self, inputs: &[Value], out: Value, ts: u32) -> Option<Value> {
         let mut res = None;
-        self.insert_and_merge(inputs, ts, false, |prev| {
+        self.insert_and_merge(inputs, ts, false, None, |prev| {
             res = prev;
             out
         });
@@ -154,6 +154,7 @@ impl Table {
         inputs: &[Value],
         ts: u32,
         subsumed: bool,
+        cost: Option<usize>,
         on_merge: impl FnOnce(Option<Value>) -> Value,
     ) {
         assert!(ts >= self.max_ts);
@@ -164,8 +165,9 @@ impl Table {
         {
             let (inp, prev) = &mut self.vals[*off];
             let prev_subsumed = prev.subsumed;
+            let prev_cost = prev.cost;
             let next = on_merge(Some(prev.value));
-            if next == prev.value && prev_subsumed == subsumed {
+            if next == prev.value && prev_subsumed == subsumed && prev_cost == cost {
                 return;
             }
             inp.stale_at = ts;
@@ -178,6 +180,8 @@ impl Table {
                     value: next,
                     timestamp: ts,
                     subsumed: subsumed || prev_subsumed,
+                    // Take miminum of cost and prev_cost, if both exist
+                    cost: cost.or(prev_cost.map(|x| cost.map_or(x, |y| x.min(y)))),
                 },
             ));
             *off = new_offset;
@@ -190,6 +194,7 @@ impl Table {
                 value: on_merge(None),
                 timestamp: ts,
                 subsumed,
+                cost: None,
             },
         ));
         let to = TableOffset {

src/serialize.rs

@@ -166,7 +166,8 @@ impl EGraph {
                 egraph_serialize::Node {
                     op: func.decl.name.to_string(),
                     eclass: class_id.clone(),
-                    cost: NotNan::new(func.decl.cost.unwrap_or(1) as f64).unwrap(),
+                    cost: NotNan::new(output.cost.unwrap_or(func.decl.cost.unwrap_or(1)) as f64)
+                        .unwrap(),
                     children,
                     subsumed: output.subsumed,
                 },

tests/eggcc-extraction.egg

@@ -1362,20 +1362,20 @@
 
 ;; if we reach a new context, union
 (rule ((= theta (Theta pred inputs outputs))
-       (= (BodyAndCost extracted cost)
+       (= (BodyAndCost extracted cost_)
           (ExtractedBody theta)))
       ((union theta extracted))
       :ruleset fast-analyses)
 (rule ((= gamma (Gamma pred inputs outputs))
-       (= (BodyAndCost extracted cost)
+       (= (BodyAndCost extracted cost_)
           (ExtractedBody gamma)))
       ((union gamma extracted))
       :ruleset fast-analyses)
 
 
 ;; if we reach the function at the top level, union
 (rule ((= func (Func name intypes outtypes body))
-       (= (VecOperandAndCost extracted cost)
+       (= (VecOperandAndCost extracted cost_)
           (ExtractedVecOperand body)))
       ((union func
               (Func name intypes outtypes extracted)))

tests/extract-costfn.egg

@@ -0,0 +1,12 @@
+(datatype E
+  (Foo String))
+
+(union (Foo "x") (Foo "y"))
+(union (Foo "y") (Foo "z"))
+
+(cost (Foo "x") 17)
+(cost (Foo "y") 11)
+(cost (Foo "z") 15)
+
+(extract (Foo "y"))
+

tests/fail-typecheck/custom-cost-wrong-type.egg

@@ -0,0 +1,3 @@
+;;; test that setting the cost of a function to a non i64 will be a type error
+(relation a ())
+(cost (a) "hi")