From 36382ccb5700428fb8ebdd3700d212cf8abe0be3 Mon Sep 17 00:00:00 2001
From: "Christoph M. Wintersteiger" <christoph@imandra.ai>
Date: Fri, 17 Nov 2023 02:28:12 +0000
Subject: [PATCH] Fix memory and concurrency issues in OCaml API (#6992)

* Fix memory and concurrency issues in OCaml API

* Undo locking changes
---
 src/api/api_datatype.cpp        | 14 ++++++++++++++
 src/api/ml/z3.ml                | 28 ++++++++++------------------
 src/api/ml/z3native_stubs.c.pre |  8 +++++++-
 src/api/z3_api.h                | 10 ++++++++++
 4 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/src/api/api_datatype.cpp b/src/api/api_datatype.cpp
index 71d1de21210..1ef4ea62688 100644
--- a/src/api/api_datatype.cpp
+++ b/src/api/api_datatype.cpp
@@ -241,6 +241,20 @@ extern "C" {
         Z3_CATCH_RETURN(nullptr);
     }
 
+    unsigned Z3_API Z3_constructor_num_fields(Z3_context c, Z3_constructor constr) {
+        Z3_TRY;
+        LOG_Z3_constructor_num_fields(c, constr);
+        RESET_ERROR_CODE();
+        mk_c(c)->reset_last_result();
+        if (!constr) {
+            SET_ERROR_CODE(Z3_INVALID_ARG, nullptr);
+            return 0;
+        }
+        constructor* c = reinterpret_cast<constructor*>(constr);
+        return c->m_field_names.size();
+        Z3_CATCH_RETURN(0);
+    }
+
 
     void Z3_API Z3_query_constructor(Z3_context c,
                                      Z3_constructor constr,
diff --git a/src/api/ml/z3.ml b/src/api/ml/z3.ml
index a5b61c8dc30..c2de2b58967 100644
--- a/src/api/ml/z3.ml
+++ b/src/api/ml/z3.ml
@@ -858,15 +858,6 @@ struct
   module Constructor =
   struct
     type constructor = Z3native.constructor
-
-    module FieldNumTable = Hashtbl.Make(struct
-        type t = AST.ast
-        let equal x y = AST.compare x y = 0
-        let hash = AST.hash
-      end)
-
-    let _field_nums = FieldNumTable.create 0
-
     let create (ctx:context) (name:Symbol.symbol) (recognizer:Symbol.symbol) (field_names:Symbol.symbol list) (sorts:Sort.sort option list) (sort_refs:int list) =
       let n = List.length field_names in
       if n <> List.length sorts then
@@ -882,10 +873,9 @@ struct
             (let f x = match x with None -> Z3native.mk_null_ast ctx | Some s -> s in
              List.map f sorts)
             sort_refs in
-        FieldNumTable.add _field_nums no n;
         no
 
-    let get_num_fields (x:constructor) = FieldNumTable.find _field_nums x
+    let get_num_fields (x:constructor) = Z3native.constructor_num_fields (gc x) x
 
     let get_constructor_decl (x:constructor) =
       let (a, _, _) = (Z3native.query_constructor (gc x) x (get_num_fields x)) in
@@ -1512,13 +1502,15 @@ struct
     in
     Z3native.apply_result_inc_ref (gc x) arn;
     let sg = Z3native.apply_result_get_num_subgoals (gc x) arn in
-    let res = if sg = 0 then
-        raise (Error "No subgoals")
-      else
-        Z3native.apply_result_get_subgoal (gc x) arn 0 in
-    Z3native.apply_result_dec_ref (gc x) arn;
-    Z3native.tactic_dec_ref (gc x) tn;
-    res
+    if sg = 0 then (
+      Z3native.apply_result_dec_ref (gc x) arn;
+      Z3native.tactic_dec_ref (gc x) tn;
+      raise (Error "No subgoals"))
+    else
+      let res:goal = Z3native.apply_result_get_subgoal (gc x) arn 0 in
+      Z3native.apply_result_dec_ref (gc x) arn;
+      Z3native.tactic_dec_ref (gc x) tn;
+      res
 
   let mk_goal = Z3native.mk_goal
 
diff --git a/src/api/ml/z3native_stubs.c.pre b/src/api/ml/z3native_stubs.c.pre
index 038b80725ee..c8afe90b909 100644
--- a/src/api/ml/z3native_stubs.c.pre
+++ b/src/api/ml/z3native_stubs.c.pre
@@ -2,6 +2,12 @@
 #include <string.h>
 #include <assert.h>
 
+#ifndef __STDC_NO_ATOMICS__
+#include <stdatomic.h>
+#else
+#define _Atomic(T) T
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -118,7 +124,7 @@ int compare_pointers(void* pt1, void* pt2) {
    blocks that get copied. */
 typedef struct {
   Z3_context ctx;
-  unsigned long obj_count;
+  _Atomic(unsigned long) obj_count;
 } Z3_context_plus_data;
 
 /* A context is wrapped to an OCaml value by storing a pointer
diff --git a/src/api/z3_api.h b/src/api/z3_api.h
index 894bd60dcc9..fafc3eebd47 100644
--- a/src/api/z3_api.h
+++ b/src/api/z3_api.h
@@ -2082,6 +2082,16 @@ extern "C" {
                                             unsigned sort_refs[]
                                             );
 
+    /**
+       \brief Retrieve the number of fields of a constructor
+
+       \param c logical context.
+       \param constr constructor.
+
+       def_API('Z3_constructor_num_fields', UINT, (_in(CONTEXT), _in(CONSTRUCTOR)))
+    */
+    unsigned Z3_API Z3_constructor_num_fields(Z3_context c, Z3_constructor constr);
+
     /**
        \brief Reclaim memory allocated to constructor.