feat: Add Datatype::update_field, FuncDecl::domain, FuncDecl::range (#455)

Co-authored-by: Mark DenHoed <mark.denhoed@cs.ox.ac.uk>

Author: willcrichton

z3/src/ast/datatype.rs

@@ -1,5 +1,5 @@
 use crate::ast::Ast;
-use crate::{Context, Sort, Symbol};
+use crate::{Context, FuncDecl, Sort, Symbol};
 use std::ffi::CString;
 use z3_sys::*;
 
@@ -35,4 +35,30 @@ impl Datatype {
             })
         }
     }
+
+    /// Update the field of the given datatype.
+    ///
+    /// The accessor should be taken from the datatype definition.
+    pub fn update_field(&self, accessor: &FuncDecl, value: &impl Ast) -> Self {
+        let c = self.ctx.get_z3_context();
+        let field_access = accessor.z3_func_decl;
+        let t = self.z3_ast;
+        let v = value.get_z3_ast();
+
+        // The Z3 docs say that the 1-indexed param should be a datatype,
+        // but that seems like a typo.
+        let domain = accessor.domain(0).expect("accessor lacks domain");
+        assert_eq!(domain, SortKind::Datatype);
+
+        let range = accessor.range();
+        assert_eq!(value.get_sort().kind(), range);
+
+        unsafe {
+            Self::wrap(
+                &self.ctx,
+                Z3_datatype_update_field(c, field_access, t, v)
+                    .expect("failed to update field of struct"),
+            )
+        }
+    }
 }

z3/src/func_decl.rs

@@ -1,7 +1,7 @@
-use std::borrow::Borrow;
 use std::convert::TryInto;
 use std::ffi::CStr;
 use std::fmt;
+use std::{borrow::Borrow, ffi::c_uint};
 use z3_sys::*;
 
 use crate::{Context, FuncDecl, Sort, Symbol, Translate, ast, ast::Ast};
@@ -247,6 +247,37 @@ impl FuncDecl {
             }
         }
     }
+
+    /// Returns the kind of the `i`-th domain (parameter) of this `FuncDecl`.
+    ///
+    /// Returns `None` if `i >= |domain|`.
+    pub fn domain(&self, i: usize) -> Option<SortKind> {
+        let z3_ctx = self.ctx.z3_ctx.0;
+        let i = c_uint::try_from(i).unwrap();
+
+        let domain_size = unsafe { Z3_get_domain_size(z3_ctx, self.z3_func_decl) };
+        if i >= domain_size {
+            return None;
+        }
+
+        Some(unsafe {
+            Z3_get_sort_kind(
+                z3_ctx,
+                Z3_get_domain(z3_ctx, self.z3_func_decl, i).expect("cannot get domain of FuncDecl"),
+            )
+        })
+    }
+
+    /// Returns the kind of range (output) of this `FuncDecl`.
+    pub fn range(&self) -> SortKind {
+        let z3_ctx = self.ctx.z3_ctx.0;
+        unsafe {
+            Z3_get_sort_kind(
+                z3_ctx,
+                Z3_get_range(z3_ctx, self.z3_func_decl).expect("cannot get range of FuncDecl"),
+            )
+        }
+    }
 }
 
 impl fmt::Display for FuncDecl {

z3/tests/lib.rs

@@ -794,6 +794,17 @@ fn test_datatype_builder() {
         .unwrap();
     solver.assert(five.eq(&five_two));
 
+    let four = ast::Int::from_i64(4);
+    let just_four = just_five
+        .as_datatype()
+        .unwrap()
+        .update_field(&maybe_int.variants[1].accessors[0], &four);
+    let four_two = maybe_int.variants[1].accessors[0]
+        .apply(&[&just_four])
+        .as_int()
+        .unwrap();
+    solver.assert(four.eq(four_two));
+
     assert_eq!(solver.check(), SatResult::Sat);
 }
 
@@ -1464,6 +1475,9 @@ fn test_ast_safe_decl() {
     assert_eq!(x_not.safe_decl().unwrap().kind(), DeclKind::NOT);
 
     let f = FuncDecl::new("f", &[&Sort::int()], &Sort::int());
+    assert_eq!(f.domain(0), Some(SortKind::Int));
+    assert_eq!(f.range(), SortKind::Int);
+
     let x = ast::Int::new_const("x");
     let f_x: ast::Int = f.apply(&[&x]).try_into().unwrap();
     let f_x_pattern: Pattern = Pattern::new(&[&f_x]);