Add new multiset primitives

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all test nits docs graphs rm-graphs doctest coverage insta-test
+.PHONY: all test nits docs graphs rm-graphs doctest coverage insta-test fixnits
 
 RUST_SRC=$(shell find . -type f -wholename '*/src/*.rs' -or -name 'Cargo.toml')
 TESTS=$(shell find tests/ -type f -name '*.egg' -not -name '*repro-*')
@@ -28,7 +28,8 @@ nits:
 fixnits:
 	@rustup component add rustfmt
 	cargo fmt
-	cargo clippy --fix --workspace --allow-dirty
+	@rustup component add rustfmt
+	cargo clippy --fix --tests --workspace --allow-dirty
 
 docs:
 	mkdir -p ${WWW}/

egglog-bridge/src/lib.rs

@@ -1216,6 +1216,7 @@ impl MergeFn {
             Primitive(_, args) => {
                 args.iter()
                     .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps));
+                write_deps.insert(egraph.uf_table);
             }
             Function(func, args) => {
                 read_deps.insert(egraph.funcs[*func].table);

src/lib.rs

@@ -37,8 +37,9 @@ pub use ast::{ResolvedExpr, ResolvedFact, ResolvedVar};
 #[cfg(feature = "bin")]
 pub use cli::*;
 use constraint::{Constraint, Problem, SimpleTypeConstraint, TypeConstraint};
+use core::CoreActionContext;
+use core::ResolvedAtomTerm;
 pub use core::{Atom, AtomTerm};
-use core::{CoreActionContext, ResolvedAtomTerm};
 pub use core::{ResolvedCall, SpecializedPrimitive};
 pub use core_relations::{BaseValue, ContainerValue, ExecutionState, Value};
 use core_relations::{ExternalFunctionId, make_external_func};
@@ -50,7 +51,7 @@ use egglog_ast::generic_ast::{Change, GenericExpr, Literal};
 use egglog_ast::span::Span;
 use egglog_ast::util::ListDisplay;
 pub use egglog_bridge::FunctionRow;
-use egglog_bridge::{ColumnTy, QueryEntry};
+use egglog_bridge::{ColumnTy, QueryEntry, UnionAction};
 use egglog_core_relations as core_relations;
 use egglog_numeric_id as numeric_id;
 use egglog_reports::{ReportLevel, RunReport};
@@ -299,6 +300,11 @@ impl Function {
     pub fn can_subsume(&self) -> bool {
         self.can_subsume
     }
+
+    /// Whether this is a let binding
+    pub fn is_let_binding(&self) -> bool {
+        self.decl.internal_let
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -350,7 +356,6 @@ impl Default for EGraph {
             proof_state,
             proof_check_program: vec![],
         };
-
         add_base_sort(&mut eg, UnitSort, span!()).unwrap();
         add_base_sort(&mut eg, StringSort, span!()).unwrap();
         add_base_sort(&mut eg, BoolSort, span!()).unwrap();
@@ -975,6 +980,28 @@ impl EGraph {
         }
     }
 
+    /// Get the list of all functions in the e-graph.
+    pub fn get_function_names(&self) -> Vec<String> {
+        self.functions.keys().cloned().collect()
+    }
+
+    /// Read the contents of the given function.
+    /// The callback f is called with each row and its subsumption status.
+    ///
+    /// Raises an error if the function does not exist.
+    pub fn function_for_each(
+        &self,
+        func_name: &str,
+        f: impl FnMut(FunctionRow<'_>),
+    ) -> Result<(), Error> {
+        let func = self
+            .functions
+            .get(func_name)
+            .ok_or_else(|| TypeError::UnboundFunction(func_name.to_string(), span!()))?;
+        self.backend.for_each(func.backend_id, f);
+        Ok(())
+    }
+
     /// Evaluates an expression, returns the sort of the expression and the evaluation result.
     pub fn eval_expr(&mut self, expr: &Expr) -> Result<(ArcSort, Value), Error> {
         let span = expr.span();
@@ -1799,6 +1826,11 @@ impl EGraph {
         self.backend
             .get_canon_repr(val, sort.column_ty(&self.backend))
     }
+
+    /// Create a new union action that can be used to union two values.
+    pub fn new_union_action(&self) -> egglog_bridge::UnionAction {
+        UnionAction::new(&self.backend)
+    }
 }
 
 struct BackendRule<'a> {

src/sort/fn.rs

@@ -15,7 +15,11 @@ use std::sync::Mutex;
 use super::*;
 
 #[derive(Clone, Debug)]
-pub struct FunctionContainer(ResolvedFunctionId, pub Vec<(ArcSort, Value)>, pub String);
+pub struct FunctionContainer(
+    pub ResolvedFunctionId,
+    pub Vec<(ArcSort, Value)>,
+    pub String,
+);
 
 // implement hash and equality based on values only not arcsorts, since
 // arcsorts are not comparable and any two values that are equal must have the same sort
@@ -57,7 +61,6 @@ impl ContainerValue for FunctionContainer {
         self.1.iter().map(|(_, v)| v).copied()
     }
 }
-
 #[derive(Debug)]
 pub struct FunctionSort {
     name: String,
@@ -161,7 +164,9 @@ impl Sort for FunctionSort {
     }
 
     fn is_eq_container_sort(&self) -> bool {
-        self.inputs.iter().any(|s| s.is_eq_sort())
+        self.inputs
+            .iter()
+            .any(|s| s.is_eq_sort() || s.is_eq_container_sort())
     }
 
     fn serialized_name(&self, container_values: &ContainerValues, value: Value) -> String {
@@ -195,6 +200,9 @@ impl Sort for FunctionSort {
             name: "unstable-app".into(),
             function: self.clone(),
         });
+
+        register_vec_primitives_for_function(eg, self.clone());
+        register_multiset_primitives_for_function(eg, self.clone());
     }
 
     fn value_type(&self) -> Option<TypeId> {

src/sort/i64.rs

@@ -41,6 +41,8 @@ impl BaseSort for I64Sort {
 
         add_literal_prim!(eg, "log2" = |a: i64| -> i64 { a.ilog2() as i64 });
 
+        add_literal_prim!(eg, "abs" = |a: i64| -?> i64 { a.checked_abs() });
+
         add_literal_prim!(eg, "<" = |a: i64, b: i64| -?> () { (a < b).then_some(()) });
         add_literal_prim!(eg, ">" = |a: i64, b: i64| -?> () { (a > b).then_some(()) });
         add_literal_prim!(eg, "<=" = |a: i64, b: i64| -?> () { (a <= b).then_some(()) });

src/sort/map.rs

@@ -91,22 +91,6 @@ impl Presort for MapSort {
                 .get_sort_by_name(v)
                 .ok_or(TypeError::UndefinedSort(v.clone(), v_span.clone()))?;
 
-            // TODO: specialize the error message
-            if k.is_eq_container_sort() {
-                return Err(TypeError::DisallowedSort(
-                    name,
-                    "Maps nested with other EqSort containers are not allowed".into(),
-                    k_span.clone(),
-                ));
-            }
-            if v.is_container_sort() {
-                return Err(TypeError::DisallowedSort(
-                    name,
-                    "Maps nested with other containers are not allowed".into(),
-                    v_span.clone(),
-                ));
-            }
-
             let out = Self {
                 name,
                 key: k.clone(),
@@ -131,7 +115,10 @@ impl ContainerSort for MapSort {
     }
 
     fn is_eq_container_sort(&self) -> bool {
-        self.key.is_eq_sort() || self.value.is_eq_sort()
+        self.key.is_eq_sort()
+            || self.value.is_eq_sort()
+            || self.key.is_eq_container_sort()
+            || self.value.is_eq_container_sort()
     }
 
     fn inner_values(
@@ -153,8 +140,8 @@ impl ContainerSort for MapSort {
         let arc = self.clone().to_arcsort();
 
         add_primitive!(eg, "map-empty" = {self.clone(): MapSort} || -> @MapContainer (arc) { MapContainer {
-            do_rebuild_keys: self.ctx.key.is_eq_sort(),
-            do_rebuild_vals: self.ctx.value.is_eq_sort(),
+            do_rebuild_keys: self.ctx.key.is_eq_sort() || self.ctx.key.is_eq_container_sort(),
+            do_rebuild_vals: self.ctx.value.is_eq_sort() || self.ctx.value.is_eq_container_sort(),
             data: BTreeMap::new()
         } });
 

src/sort/mod.rs

@@ -76,7 +76,7 @@ pub trait Sort: Any + Send + Sync + Debug {
         false
     }
 
-    // return true if it is a container sort that contains ids.
+    // return true if it is a container sort that contains ids or other container sorts that contain ids.
     // only eq_sort and eq_container_sort need to be canonicalized.
     fn is_eq_container_sort(&self) -> bool {
         false