From bcf567247f3b3a585c0cedf2a592dca350f398d6 Mon Sep 17 00:00:00 2001
From: Ryan Scott <rscott@galois.com>
Date: Wed, 27 Nov 2024 13:48:18 -0500
Subject: [PATCH] Add prime to isNumeric, define cryPrime SMT-LIB function

Fixes #1777.
---
 lib/CryptolTC.z3                    | 17 +++++++++++++++++
 src/Cryptol/TypeCheck/Solver/SMT.hs |  4 +++-
 src/Cryptol/TypeCheck/TypePat.hs    |  5 ++++-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/lib/CryptolTC.z3 b/lib/CryptolTC.z3
index 391a85151..f0a0ecba1 100644
--- a/lib/CryptolTC.z3
+++ b/lib/CryptolTC.z3
@@ -75,6 +75,23 @@
   (cryBool true)
 )
 
+
+(declare-fun cryPrimeUnknown (Int) Bool)
+
+(define-fun cryPrimeTable ((x Int)) Bool
+  (or (= x 2)
+  (or (= x 3)
+  (or (= x 5)
+  (or (= x 7)
+  (cryPrimeUnknown x)))))
+)
+
+(define-fun cryPrime ((x InfNat)) MaybeBool
+  (ite (isErr x) cryErrProp
+    (cryBool (and (isFin x) (cryPrimeTable (value x)))))
+)
+
+
 ; ------------------------------------------------------------------------------
 ; Basic Cryptol assume/assert
 
diff --git a/src/Cryptol/TypeCheck/Solver/SMT.hs b/src/Cryptol/TypeCheck/Solver/SMT.hs
index 99f1aabb4..7cba94637 100644
--- a/src/Cryptol/TypeCheck/Solver/SMT.hs
+++ b/src/Cryptol/TypeCheck/Solver/SMT.hs
@@ -348,7 +348,8 @@ flatGoal g = [ g { goal = p } | p <- pSplitAnd (goal g) ]
 
 -- | Assumes no 'And'
 isNumeric :: Prop -> Bool
-isNumeric ty = matchDefault False $ msum [ is (|=|), is (|/=|), is (|>=|), is aFin ]
+isNumeric ty = matchDefault False $ msum [ is (|=|), is (|/=|), is (|>=|)
+                                         , is aFin, is aPrime ]
   where
   is f = f ty >> return True
 
@@ -368,6 +369,7 @@ toSMT tvs ty = matchDefault (panic "toSMT" [ "Unexpected type", show ty ])
   , aNat            ~> "cryNat"
 
   , aFin            ~> "cryFin"
+  , aPrime          ~> "cryPrime"
   , (|=|)           ~> "cryEq"
   , (|/=|)          ~> "cryNeq"
   , (|>=|)          ~> "cryGeq"
diff --git a/src/Cryptol/TypeCheck/TypePat.hs b/src/Cryptol/TypeCheck/TypePat.hs
index 53023830c..c95ee9f0f 100644
--- a/src/Cryptol/TypeCheck/TypePat.hs
+++ b/src/Cryptol/TypeCheck/TypePat.hs
@@ -24,7 +24,7 @@ module Cryptol.TypeCheck.TypePat
   , aRec
   , (|->|)
 
-  , aFin, (|=|), (|/=|), (|>=|)
+  , aFin, aPrime, (|=|), (|/=|), (|>=|)
   , aAnd
   , aTrue
 
@@ -165,6 +165,9 @@ aRec = \a -> case tNoUser a of
 aFin :: Pat Prop Type
 aFin = tp PFin ar1
 
+aPrime :: Pat Prop Type
+aPrime = tp PPrime ar1
+
 (|=|) :: Pat Prop (Type,Type)
 (|=|) = tp PEqual ar2