From 9a038e1889732025fe3f5b65440bb822bfb3e6d3 Mon Sep 17 00:00:00 2001
From: Ryan <mccleeary@galois.com>
Date: Mon, 29 Jul 2024 16:08:04 -0600
Subject: [PATCH] Enhance TypeCheck to reason more about exponents.

Move (x >= 2, x^a + x^b = x^c => a + b = c) check from Improve.hs to Numeric.hs.
Add (Nat a, a * a^x = a^y => 1 + x = y) to SimpType.hs
Add (x >= 2,  x^a >= x^b => a >= b) to Numeric.hs
---
 src/Cryptol/TypeCheck/SimpType.hs       |  5 +++-
 src/Cryptol/TypeCheck/Solver/Improve.hs | 36 ++++++++-----------------
 src/Cryptol/TypeCheck/Solver/Numeric.hs | 28 +++++++++++++++++++
 3 files changed, 43 insertions(+), 26 deletions(-)

diff --git a/src/Cryptol/TypeCheck/SimpType.hs b/src/Cryptol/TypeCheck/SimpType.hs
index af3317f62..971da0703 100644
--- a/src/Cryptol/TypeCheck/SimpType.hs
+++ b/src/Cryptol/TypeCheck/SimpType.hs
@@ -154,7 +154,10 @@ tMul x y
            , Just b' <- tIsNum b
            -- XXX: similar for a = b * k?
            , n == b' = tSub a (tMod a b)
-
+           -- c * c ^ x = c ^ (1 + x)
+           | TCon (TF TCExp) [a,b] <- t'
+           , Just n' <- tIsNum a
+           , n == n' = tf2 TCExp a (tAdd (tNum (1::Int)) b)
 
            | otherwise = tf2 TCMul (tNum n) t
     where t' = tNoUser t
diff --git a/src/Cryptol/TypeCheck/Solver/Improve.hs b/src/Cryptol/TypeCheck/Solver/Improve.hs
index 61d4cc41d..9114a656b 100644
--- a/src/Cryptol/TypeCheck/Solver/Improve.hs
+++ b/src/Cryptol/TypeCheck/Solver/Improve.hs
@@ -15,7 +15,6 @@ import Cryptol.TypeCheck.Type
 import Cryptol.TypeCheck.SimpType as Mk
 import Cryptol.TypeCheck.Solver.Types
 import Cryptol.TypeCheck.Solver.Numeric.Interval
-import Cryptol.TypeCheck.Solver.InfNat(Nat'(..))
 import Cryptol.TypeCheck.TypePat
 import Cryptol.TypeCheck.Subst
 
@@ -69,11 +68,10 @@ improveLit impSkol prop =
 improveEq :: Bool -> Ctxt -> Prop -> Match (Subst,[Prop])
 improveEq impSkol fins prop =
   do (lhs,rhs) <- (|=|) prop
-     rewriteNotSym lhs rhs <|> rewriteNotSym rhs lhs
-                           <|> rewriteSym rhs lhs
+     rewrite lhs rhs <|> rewrite rhs lhs
   where
-  rewriteNotSym this other =
-
+  rewrite this other =
+    
     -- x == t       ~>   (x := t)
     do x <- aTVar this
        guard (considerVar x)
@@ -83,27 +81,15 @@ improveEq impSkol fins prop =
     <|>
 
     -- x + y  = t     ~> (x := t - y, y <= t)
-    do (v,s) <- isSum this
-       case singleSubst v (Mk.tSub other s) of
+    do (x,y) <- isSum this
+       case singleSubst x (Mk.tSub other y) of
          Left _ -> mzero
-         Right su -> return (su, [ other >== s ])
-
-
-  rewriteSym this other =
-
-    -- (a >= 2, fin a, a ^ x == a ^ y) -> x = y
-    do (a,x) <- matches this  ((|^|), __, __)
-       (b,y) <- matches other ((|^|), __, __)
-       guard (a == b)
-       let i = intervalFor a
-       guard (iLower i >= Nat 2 && iIsFin i)
-       improveEq impSkol fins (x =#= y)
-
-
-  isSum t = do (v,s) <- matches t (anAdd, aTVar, __)
-               valid v s
-        <|> do (s,v) <- matches t (anAdd, __, aTVar)
-               valid v s
+         Right su -> return (su, [ other >== y ])
+  
+  isSum t  = do (v,s) <- matches t (anAdd, aTVar, __)
+                valid v s
+         <|> do (s,v) <- matches t (anAdd, __, aTVar)
+                valid v s
 
   valid v s = do let i = intervalFor s
                  guard (considerVar v && v `Set.notMember` fvs s && iIsFin i)
diff --git a/src/Cryptol/TypeCheck/Solver/Numeric.hs b/src/Cryptol/TypeCheck/Solver/Numeric.hs
index 60123183a..db47893b7 100644
--- a/src/Cryptol/TypeCheck/Solver/Numeric.hs
+++ b/src/Cryptol/TypeCheck/Solver/Numeric.hs
@@ -29,6 +29,7 @@ import Cryptol.TypeCheck.SimpType as Simp
 
 
 -- | Try to solve @t1 = t2@
+-- We should move EXP rule to here.
 cryIsEqual :: Ctxt -> Type -> Type -> Solved
 cryIsEqual ctxt t1 t2 =
   matchDefault Unsolved $
@@ -38,6 +39,7 @@ cryIsEqual ctxt t1 t2 =
     <|> (aTVar t1 >>= tryEqVar t2)
     <|> (aTVar t2 >>= tryEqVar t1)
     <|> ( guard (t1 == t2) >> return (SolvedIf []))
+    <|> tryEqExp t1 t2
     <|> tryEqMin t1 t2
     <|> tryEqMin t2 t1
     <|> tryEqMins t1 t2
@@ -61,6 +63,7 @@ cryIsGeq i t1 t2 =
     <|> (aNat' t1 >>= tryGeqKThan i t2)
     <|> (aNat' t2 >>= tryGeqThanK i t1)
     <|> (aTVar t2 >>= tryGeqThanVar i t1)
+    <|> tryGeqExp i t1 t2
     <|> tryGeqThanSub i t1 t2
     <|> (geqByInterval i t1 t2)
     <|> (guard (t1 == t2) >> return (SolvedIf []))
@@ -137,6 +140,17 @@ tryGeqThanK _ t (Nat k) =
   -- XXX: K1 ^^ n >= K2
 
 
+-- (x >= 2 && x^a >= x^b) => a >= b
+tryGeqExp :: Ctxt -> Type -> Type -> Match Solved
+tryGeqExp _ x y = 
+      do  (x_1, a) <- (|^|) x
+          (x_2, b) <- (|^|) y
+          guard (x_1 == x_2)
+          n <- aNat x_1
+          guard (n >= 2)
+          return $ SolvedIf [ a >== b ]
+
+
 tryGeqThanSub :: Ctxt -> Type -> Type -> Match Solved
 tryGeqThanSub _ x y =
 
@@ -223,6 +237,20 @@ tryCancelVar ctxt p t1 t2 =
 
 
 
+-- if (x >= 2) && x^a*x^b = x^c => a+b = c
+tryEqExp :: Type -> Type -> Match Solved
+tryEqExp x y = check x y <|> check y x
+  where 
+    check i j =
+      do  (m1,m2) <- aMul i
+          (x_1, a) <- (|^|) m1
+          (x_2, b) <- (|^|) m2
+          (x_3, c) <- (|^|) j
+          guard (x_1 == x_2 && x_2 == x_3)
+          n <- aNat x_1
+          guard (n >= 2)
+          return $ SolvedIf [ tAdd a b =#= c ]
+  
 -- min t1 t2 = t1 ~> t1 <= t2
 tryEqMin :: Type -> Type -> Match Solved
 tryEqMin x y =