Merge pull request #239 from rtjoa/vec-set

Add `vec-set`

src/sort/vec.rs

@@ -27,6 +27,7 @@ impl VecSort {
             "vec-contains".into(),
             "vec-length".into(),
             "vec-get".into(),
+            "vec-set".into(),
         ]
     }
 
@@ -138,7 +139,12 @@ impl Sort for VecSort {
         });
         typeinfo.add_primitive(Get {
             name: "vec-get".into(),
-            vec: self,
+            vec: self.clone(),
+            i64: typeinfo.get_sort(),
+        });
+        typeinfo.add_primitive(Set {
+            name: "vec-set".into(),
+            vec: self.clone(),
             i64: typeinfo.get_sort(),
         })
     }
@@ -422,3 +428,35 @@ impl PrimitiveLike for Get {
         vec.get(index as usize).copied()
     }
 }
+
+struct Set {
+    name: Symbol,
+    vec: Arc<VecSort>,
+    i64: Arc<I64Sort>,
+}
+
+impl PrimitiveLike for Set {
+    fn name(&self) -> Symbol {
+        self.name
+    }
+
+    fn accept(&self, types: &[ArcSort]) -> Option<ArcSort> {
+        match types {
+            [vec, index, el]
+                if vec.name() == self.vec.name
+                    && index.name() == "i64".into()
+                    && el.name() == self.vec.element_name() =>
+            {
+                Some(self.vec.clone())
+            }
+            _ => None,
+        }
+    }
+
+    fn apply(&self, values: &[Value]) -> Option<Value> {
+        let mut vec = ValueVec::load(&self.vec, &values[0]);
+        let index = i64::load(&self.i64, &values[1]);
+        vec[index as usize] = values[2];
+        vec.store(&self.vec)
+    }
+}

tests/vec.egg

@@ -20,3 +20,6 @@
 
 ; Test vec-get
 (check (= (vec-get (vec-of 1 2 3) 1) 2))
+
+; Test vec-set
+(check (= (vec-set (vec-of 1 2 3) 1 4) (vec-of 1 4 3)))