Add multiset-sum primitive

src/sort/multiset.rs

@@ -95,6 +95,14 @@ mod inner {
             }
             multiset
         }
+
+        /// Compute the sum of two multisets.
+        pub fn sum(self, MultiSet(other_map, other_count): Self) -> Self {
+            Self(
+                self.0.union_with(other_map, std::ops::Add::add),
+                self.1 + other_count,
+            )
+        }
     }
 }
 
@@ -130,6 +138,7 @@ impl Presort for MultiSetSort {
             "multiset-not-contains".into(),
             "multiset-remove".into(),
             "multiset-length".into(),
+            "multiset-sum".into(),
             "unstable-multiset-map".into(),
         ]
     }
@@ -233,6 +242,10 @@ impl Sort for MultiSetSort {
             name: "multiset-pick".into(),
             multiset: self.clone(),
         });
+        typeinfo.add_primitive(Sum {
+            name: "multiset-sum".into(),
+            multiset: self.clone(),
+        });
         let inner_name = self.element.name();
         let fn_sort = typeinfo.get_sort_by(|s: &Arc<FunctionSort>| {
             (s.output.name() == inner_name)
@@ -535,6 +548,41 @@ impl PrimitiveLike for Pick {
     }
 }
 
+struct Sum {
+    name: Symbol,
+    multiset: Arc<MultiSetSort>,
+}
+
+impl PrimitiveLike for Sum {
+    fn name(&self) -> Symbol {
+        self.name
+    }
+
+    fn get_type_constraints(&self, span: &Span) -> Box<dyn TypeConstraint> {
+        SimpleTypeConstraint::new(
+            self.name(),
+            vec![
+                self.multiset.clone(),
+                self.multiset.clone(),
+                self.multiset.clone(),
+            ],
+            span.clone(),
+        )
+        .into_box()
+    }
+
+    fn apply(
+        &self,
+        values: &[Value],
+        _sorts: (&[ArcSort], &ArcSort),
+        _egraph: Option<&mut EGraph>,
+    ) -> Option<Value> {
+        let lhs_multiset = ValueMultiSet::load(&self.multiset, &values[0]);
+        let rhs_multiset = ValueMultiSet::load(&self.multiset, &values[1]);
+        lhs_multiset.sum(rhs_multiset).store(&self.multiset)
+    }
+}
+
 struct Map {
     name: Symbol,
     multiset: Arc<MultiSetSort>,

tests/multiset.egg

@@ -56,3 +56,15 @@
     (multiset-of (Num 1) (Num 4) (Num 9))
     squared-xs
 ))
+
+;; sum
+(check (=
+    (multiset-sum (multiset-of (Num 1) (Num 2) (Num 3)) (multiset-of (Num 1) (Num 2) (Num 4)))
+    (multiset-of (Num 1) (Num 4) (Num 2) (Num 3) (Num 2) (Num 1))
+))
+
+;; verify that sum computes length
+(check (=
+    (multiset-length (multiset-sum (multiset-of (Num 1) (Num 2) (Num 3)) (multiset-of (Num 1) (Num 2) (Num 4))))
+    6
+))