From 0dc7f3e7c83f128bca0c6e688d727d56fcc187b3 Mon Sep 17 00:00:00 2001
From: Romain Ruetschi <romain.ruetschi@gmail.com>
Date: Wed, 4 Sep 2019 15:26:21 +0200
Subject: [PATCH] Add BigInt extraction and test for Int->BigInt conversion

---
 src/main/scala/z3/scala/Z3Model.scala       | 32 +++++++++++----------
 src/test/scala/z3/scala/IntArith.scala      | 29 +++++++++++++++++++
 src/test/scala/z3/scala/IntExtraction.scala | 26 +++++++++++++++++
 3 files changed, 72 insertions(+), 15 deletions(-)
 create mode 100644 src/test/scala/z3/scala/IntExtraction.scala

diff --git a/src/main/scala/z3/scala/Z3Model.scala b/src/main/scala/z3/scala/Z3Model.scala
index f50e675..60e3dd2 100644
--- a/src/main/scala/z3/scala/Z3Model.scala
+++ b/src/main/scala/z3/scala/Z3Model.scala
@@ -3,28 +3,30 @@ package z3.scala
 import com.microsoft.z3.Native
 
 object Z3Model {
+  implicit def ast2bigint(model: Z3Model, ast: Z3AST): Option[BigInt] = {
+    model.eval(ast) flatMap { res =>
+      model.context.getNumeralInt(res).value.map(BigInt(_)).orElse {
+        Some(BigInt(res.toString, 10))
+      }
+    }
+  }
+
   implicit def ast2int(model: Z3Model, ast: Z3AST): Option[Int] = {
-    val res = model.eval(ast)
-    if (res.isEmpty)
-      None
-    else
-      model.context.getNumeralInt(res.get).value
+    model.eval(ast) flatMap { res =>
+      model.context.getNumeralInt(res).value
+    }
   }
 
   implicit def ast2bool(model: Z3Model, ast: Z3AST): Option[Boolean] = {
-    val res = model.eval(ast)
-    if (res.isEmpty)
-      None
-    else
-      model.context.getBoolValue(res.get)
+    model.eval(ast) flatMap { res =>
+      model.context.getBoolValue(res)
+    }
   }
 
   implicit def ast2char(model: Z3Model, ast: Z3AST): Option[Char] = {
-    val res = model.eval(ast)
-    if (res.isEmpty)
-      None
-    else
-      model.context.getNumeralInt(res.get).value.map(_.toChar)
+    model.eval(ast) flatMap { res =>
+      model.context.getNumeralInt(res).value.map(_.toChar)
+    }
   }
 }
 
diff --git a/src/test/scala/z3/scala/IntArith.scala b/src/test/scala/z3/scala/IntArith.scala
index e6c7d2d..c9956b7 100644
--- a/src/test/scala/z3/scala/IntArith.scala
+++ b/src/test/scala/z3/scala/IntArith.scala
@@ -40,5 +40,34 @@ class IntArith extends FunSuite with Matchers {
 
     z3.delete
   }
+
+  testIntToBigIntConversion(0)
+  testIntToBigIntConversion(42)
+  testIntToBigIntConversion(-42)
+  testIntToBigIntConversion(Int.MinValue)
+  testIntToBigIntConversion(Int.MaxValue)
+
+  private def testIntToBigIntConversion(value: Int): Unit = test(s"Integer conversion: $value") {
+    val z3 = new Z3Context("MODEL" -> true)
+
+    val Int  = z3.mkIntSort
+    val BV32 = z3.mkBVSort(32)
+
+    val in  = z3.mkConst(z3.mkStringSymbol("in"), BV32)
+    val out = z3.mkConst(z3.mkStringSymbol("out"), Int)
+
+    val solver = z3.mkSolver
+
+    solver.assertCnstr(z3.mkEq(in,  z3.mkInt(value, BV32)))
+    solver.assertCnstr(z3.mkEq(out, z3.mkBV2Int(in, true)))
+    solver.assertCnstr(z3.mkEq(out, z3.mkInt(value, Int)))
+
+    val (sol, model) = solver.checkAndGetModel
+
+    sol should equal(Some(true))
+    model.evalAs[BigInt](out) should equal(Some(BigInt(value)))
+
+    z3.delete
+  }
 }
 
diff --git a/src/test/scala/z3/scala/IntExtraction.scala b/src/test/scala/z3/scala/IntExtraction.scala
new file mode 100644
index 0000000..cfdccfd
--- /dev/null
+++ b/src/test/scala/z3/scala/IntExtraction.scala
@@ -0,0 +1,26 @@
+package z3.scala
+
+import org.scalatest.{FunSuite, Matchers}
+
+class IntExtraction extends FunSuite with Matchers {
+
+  test(s"BigInt extraction") {
+    val z3 = new Z3Context("MODEL" -> true)
+
+    val i = z3.mkIntSort
+    val x = z3.mkConst(z3.mkStringSymbol("x"), i)
+    val m = z3.mkInt(Int.MaxValue, i)
+
+    val solver = z3.mkSolver
+
+    solver.assertCnstr(z3.mkEq(x, z3.mkAdd(m, m)))
+
+    val (sol, model) = solver.checkAndGetModel
+
+    sol should equal(Some(true))
+    model.evalAs[BigInt](x) should equal(Some(BigInt(Int.MaxValue) * 2))
+
+    z3.delete
+  }
+}
+