From fc6c4c98e26fc0bc1c5c50e1376ff9201f94b880 Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Thu, 9 May 2024 14:52:49 -0700
Subject: [PATCH] initial warppers for seq-map/seq-fold

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
---
 src/api/api_ast.cpp     |  4 ++++
 src/api/api_seq.cpp     |  5 +++++
 src/api/api_util.h      | 17 +++++++++++++++++
 src/api/python/z3/z3.py | 13 +++++++++++++
 src/api/z3_api.h        | 28 ++++++++++++++++++++++++++++
 5 files changed, 67 insertions(+)

diff --git a/src/api/api_ast.cpp b/src/api/api_ast.cpp
index 424b361f3..931167110 100644
--- a/src/api/api_ast.cpp
+++ b/src/api/api_ast.cpp
@@ -1310,6 +1310,10 @@ extern "C" {
             case OP_SEQ_INDEX: return Z3_OP_SEQ_INDEX;
             case OP_SEQ_TO_RE: return Z3_OP_SEQ_TO_RE;
             case OP_SEQ_IN_RE: return Z3_OP_SEQ_IN_RE;
+            case OP_SEQ_MAP: return Z3_OP_SEQ_MAP;
+            case OP_SEQ_MAPI: return Z3_OP_SEQ_MAPI;
+            case OP_SEQ_FOLDL: return Z3_OP_SEQ_FOLDL;
+            case OP_SEQ_FOLDLI: return Z3_OP_SEQ_FOLDLI;
 
             case _OP_STRING_STRREPL: return Z3_OP_SEQ_REPLACE;
             case _OP_STRING_CONCAT: return Z3_OP_SEQ_CONCAT;
diff --git a/src/api/api_seq.cpp b/src/api/api_seq.cpp
index 6a9d0f81c..2b87ef290 100644
--- a/src/api/api_seq.cpp
+++ b/src/api/api_seq.cpp
@@ -348,5 +348,10 @@ extern "C" {
     MK_UNARY(Z3_mk_char_from_bv, mk_c(c)->get_char_fid(), OP_CHAR_FROM_BV, SKIP);
     MK_UNARY(Z3_mk_char_is_digit, mk_c(c)->get_char_fid(), OP_CHAR_IS_DIGIT, SKIP);
 
+    MK_BINARY(Z3_mk_seq_map, mk_c(c)->get_seq_fid(), OP_SEQ_MAP, SKIP); 
+    MK_TERNARY(Z3_mk_seq_mapi, mk_c(c)->get_seq_fid(), OP_SEQ_MAPI, SKIP);
+    MK_TERNARY(Z3_mk_seq_foldl, mk_c(c)->get_seq_fid(), OP_SEQ_FOLDL, SKIP);
+    MK_FOURARY(Z3_mk_seq_foldli, mk_c(c)->get_seq_fid(), OP_SEQ_FOLDLI, SKIP);
+
 
 };
diff --git a/src/api/api_util.h b/src/api/api_util.h
index 174d75144..7b16229f4 100644
--- a/src/api/api_util.h
+++ b/src/api/api_util.h
@@ -160,6 +160,23 @@ Z3_ast Z3_API NAME(Z3_context c, Z3_ast n1, Z3_ast n2) {        \
     MK_TERNARY_BODY(NAME, FID, OP, EXTRA_CODE);                          \
 }
 
+#define MK_FOURARY_BODY(NAME, FID, OP, EXTRA_CODE)               \
+    Z3_TRY;                                                     \
+    RESET_ERROR_CODE();                                         \
+    EXTRA_CODE;                                                 \
+    expr * args[4] = { to_expr(n1), to_expr(n2), to_expr(n3), to_expr(n4) }; \
+    ast* a = mk_c(c)->m().mk_app(FID, OP, 0, 0, 4, args);       \
+    mk_c(c)->save_ast_trail(a);                                 \
+    check_sorts(c, a);                                          \
+    RETURN_Z3(of_ast(a));                                       \
+    Z3_CATCH_RETURN(0);
+
+#define MK_FOURARY(NAME, FID, OP, EXTRA_CODE)                            \
+    Z3_ast Z3_API NAME(Z3_context c, Z3_ast n1, Z3_ast n2, Z3_ast n3, Z3_ast n4) { \
+    LOG_ ## NAME(c, n1, n2, n3, n4);                                        \
+    MK_FOURARY_BODY(NAME, FID, OP, EXTRA_CODE);                          \
+}
+
 #define MK_NARY(NAME, FID, OP, EXTRA_CODE)                              \
 Z3_ast Z3_API NAME(Z3_context c, unsigned num_args, Z3_ast const* args) { \
     Z3_TRY;                                                             \
diff --git a/src/api/python/z3/z3.py b/src/api/python/z3/z3.py
index 9a3dadda2..44a4aba84 100644
--- a/src/api/python/z3/z3.py
+++ b/src/api/python/z3/z3.py
@@ -11210,6 +11210,19 @@ def Length(s):
     s = _coerce_seq(s)
     return ArithRef(Z3_mk_seq_length(s.ctx_ref(), s.as_ast()), s.ctx)
 
+def SeqMap(f, s):
+    """Map function 'f' over sequence 's'"""
+    ctx = _get_ctx2(f, s)
+    s = _coerce_seq(s, ctx)
+    return _to_expr_ref(Z3_mk_seq_map(s.ctx_ref(), f.as_ast(), s.as_ast()), ctx)
+
+def SeqMapI(f, i, s):
+    """Map function 'f' over sequence 's' at index 'i'"""
+    ctx = _get_ctx(f, s)
+    s = _coerce_seq(s, ctx)
+    if not is_expr(i):
+        i = _py2expr(i)
+    return _to_expr_ref(Z3_mk_seq_mapi(s.ctx_ref(), f.as_ast(), i.as_ast(), s.as_ast()), ctx)
 
 def StrToInt(s):
     """Convert string expression to integer
diff --git a/src/api/z3_api.h b/src/api/z3_api.h
index cbf9803db..312acc268 100644
--- a/src/api/z3_api.h
+++ b/src/api/z3_api.h
@@ -1193,6 +1193,10 @@ typedef enum {
     Z3_OP_SEQ_LAST_INDEX,
     Z3_OP_SEQ_TO_RE,
     Z3_OP_SEQ_IN_RE,
+    Z3_OP_SEQ_MAP,            
+    Z3_OP_SEQ_MAPI,           
+    Z3_OP_SEQ_FOLDL,          
+    Z3_OP_SEQ_FOLDLI,         
 
     // strings
     Z3_OP_STR_TO_INT,
@@ -3798,6 +3802,30 @@ extern "C" {
     */
     Z3_ast Z3_API Z3_mk_seq_last_index(Z3_context c, Z3_ast s, Z3_ast substr);
 
+    /**
+      \brief Create a map of the function \c f over the sequence \c s.
+      def_API('Z3_mk_seq_map', AST ,(_in(CONTEXT), _in(AST), _in(AST)))
+    */
+    Z3_ast Z3_API Z3_mk_seq_map(Z3_context c, Z3_ast f, Z3_ast s);
+
+    /**
+     \brief Create a map of the function \c f over the sequence \c s starting at index \c i.
+     def_API('Z3_mk_seq_mapi', AST ,(_in(CONTEXT), _in(AST), _in(AST), _in(AST)))
+    */
+    Z3_ast Z3_API Z3_mk_seq_mapi(Z3_context c, Z3_ast f, Z3_ast i, Z3_ast s);
+
+    /**
+      \brief Create a fold of the function \c f over the sequence \c s with accumulator a.
+      def_API('Z3_mk_seq_foldl', AST ,(_in(CONTEXT), _in(AST), _in(AST), _in(AST)))
+    */
+    Z3_ast Z3_API Z3_mk_seq_foldl(Z3_context c, Z3_ast f, Z3_ast a, Z3_ast s);
+
+    /**
+       \brief Create a fold with index tracking of the function \c f over the sequence \c s with accumulator \c a starting at index \c i.
+       def_API('Z3_mk_seq_foldli', AST ,(_in(CONTEXT), _in(AST), _in(AST), _in(AST), _in(AST)))
+    */
+    Z3_ast Z3_API Z3_mk_seq_foldli(Z3_context c, Z3_ast f, Z3_ast i, Z3_ast a, Z3_ast s);
+
     /**
        \brief Convert string to integer.