Main.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

-- | This tutorial demonstrates Grisette's core capabilities through a simple
-- expression language.
--
-- What you'll learn:
-- 1. How to define the syntax of the language using GADTs
-- 2. How to define the semantics of the language
-- 3. How to verify expression equivalence using SMT solvers
-- 4. How to synthesize equivalent expressions using CEGIS
--
-- Key concepts:
--
-- * Symbolic values (`SymInteger`, `SymBool`)
-- * Union types for representing choices
-- * Interpreter construction
-- * SMT-based verification
-- * Program synthesis
--
-- Prerequisites: Basic Haskell, familiarity with GADTs
module Main (main) where

import Grisette

-- * The expression language

-- | In this example, we will define the syntax and semantics of a simple
-- boolean and integer expression language.
--
-- The language syntax follows the following grammar:
--
-- Expr -> IntExpr | BoolExpr
-- IntExpr -> IntVal int
--          | Add IntExpr IntExpr
--          | Mul IntExpr IntExpr
-- BoolExpr -> BoolVal bool
--           | BAnd BoolExpr BoolExpr
--           | BOr BoolExpr BoolExpr
--           | Eq Expr Expr

-- * Build the syntax of the expression language.

-- A symbolic expression can be represented in Grisette as a GADT as follows.
--
-- We will see how to define values of this expression type later.
data Expr a where
  -- | A symbolic integer value.
  --
  -- The symbolic integer (`SymInteger`) type is a symbolic (primitive) type,
  -- representing SMT integer terms.
  IntVal :: SymInteger -> IntExpr
  -- | A symbolic boolean value.
  --
  -- `SymBool` is another symbolic (primitive) type, representing SMT boolean
  -- terms.
  BoolVal :: SymBool -> BoolExpr
  -- | Addition of two symbolic integer expressions.
  --
  -- Here, `IntUExpr` is a type alias for `UExpr SymInteger`, which is a union
  -- of `Expr SymInteger`.
  --
  -- A union value is a value representing a choice of symbolic expressions.
  -- Our addition operator can take choices as operands to represent a program
  -- space.
  Add :: IntUExpr -> IntUExpr -> IntExpr
  -- | Multiplication of two symbolic integer expressions.
  Mul :: IntUExpr -> IntUExpr -> IntExpr
  -- | Logical AND of two symbolic boolean expressions.
  BAnd :: BoolUExpr -> BoolUExpr -> BoolExpr
  -- | Logical OR of two symbolic boolean expressions.
  BOr :: BoolUExpr -> BoolUExpr -> BoolExpr
  -- | Equality check between two symbolic expressions.
  --
  -- `BasicSymPrim` is a constraint that contains all the symbolic primitive
  -- types that Grisette supports, including `SymInteger` and `SymBool`.
  Eq :: (BasicSymPrim a) => UExpr a -> UExpr a -> BoolExpr

-- Some type aliases for clarity.
type IntExpr = Expr SymInteger

type BoolExpr = Expr SymBool

type UExpr a = Union (Expr a)

type IntUExpr = UExpr SymInteger

type BoolUExpr = UExpr SymBool

-- The `derive` and `makeSmartCtor` functions are provided by Grisette
-- to automatically derive some instances and construct some smart constructors
-- for the GADT.
--
-- The GADT will then be fully compatible with Grisette library. And we will see
-- their usage in the following code.
derive [''Expr] basicClasses0
makeSmartCtor ''Expr

-- Here we construct some symbolic expressions.
-- We use the data constructors to construct the expressions.
e0 = IntVal 1

e1 = IntVal "a"

-- |
-- >>> target = Add (intVal "a") (intVal "a")
-- >>> sketch = Mul (intVal "a") (intVal "c")
--
-- >>> cegisForAll z3 target $ cegisPostCond $ eval target .== eval sketch
-- ([],CEGISSuccess (Model {c -> 2 :: Integer}))
e2 = IntVal ("a" + 1)

-- intVal is a smart constructor provided by Grisette to construct `Union` of
-- expressions.
e3 = intVal 1

e4 = Add (intVal "a") e3

add2 = add (intVal "a") (intVal 2)

mul2 = mul (intVal "a") (intVal 2)

-- We can then use `mrgIf` function to perform a choice between two symbolic
-- expressions.
--
-- The following code will construct the following expression:
--
-- >>> chosen
-- {If choice (Add {IntVal a} {IntVal 2}) (Mul {IntVal a} {IntVal 2})}
--
-- In this code, the `choice` can be used to choose between `add2`(i.e., a + 2)
-- and `mul2` (i.e., a * 2).
--
-- A synthesizer can then pick true or false for the `choice` variable to
-- decide which expression to pick. If the synthesizer picks true, the result is

-- $a+2$; otherwise, it is $a*2$.

chosen :: IntUExpr
chosen = mrgIf "choice" add2 mul2

-- * Defining the semantics of the expression language.

-- | Evaluation of a symbolic expression.
--
-- >>> eval (IntVal 1)
-- 1
-- >>> eval (IntVal "a")
-- a
-- >>> eval (Add (intVal 1) (intVal 2))
-- 3
-- >>> eval (Eq (intVal "x") (intVal 2))
-- (= x 2)
--
-- The `.#` combinator handles the application on `Union`s:
--
-- 1. It extracts all the symbolic choices from the expression.
-- 2. It applies the `eval` function to each of the symbolic choices.
-- 3. It merges the results into a single value.
--
-- as shown in the following pseudo code:
--
-- eval {If a choice1 choice2) = {If a (eval choice1) (eval choice2)}
--
-- >>> expr = mrgIf "choice" (add (intVal "a") (intVal 2)) (mul (intVal "a") (intVal 2)) :: IntUExpr
-- >>> eval .# expr
eval :: Expr a -> a
eval (IntVal a) = a
eval (BoolVal a) = a
eval (Add a b) = eval .# a + eval .# b
eval (Mul a b) = eval .# a * eval .# b
eval (BAnd a b) = eval .# a .&& eval .# b
eval (BOr a b) = eval .# a .|| eval .# b
eval (Eq a b) = eval .# a .== eval .# b

-- With the syntax and semantics defined, it is easy to build solver-aided tools
-- in Grisette. The following code builds a verifier to check if two expressions
-- are equivalent.
verifyEquivalent :: (BasicSymPrim a) => Expr a -> Expr a -> IO ()
verifyEquivalent e1 e2 = do
  -- We check if there exists a counter-example that falsifies the equivalence
  -- of the two expressions.
  res <- solve z3 $ eval e1 ./= eval e2
  case res of
    -- If the solver cannot find such a counter-example, the two expressions are
    -- equivalent.
    Left Unsat -> putStrLn "The two expressions are equivalent"
    -- If the solver returns an unexpected response, we print the error.
    Left err -> putStrLn $ "The solver returned unexpected response: " <> show err
    -- If the solver finds a counter-example, we print the counter-example and
    -- the two expressions.
    Right model -> do
      putStrLn "The two expressions are not equivalent, under the model:"
      print model
      putStrLn $ "lhs: " <> show e1
      putStrLn $ "rhs: " <> show e2
      -- We can use tools like `evalSym` to evaluate the expressions under the
      -- model to see why the two expressions are not equivalent.
      putStrLn $ "lhs evaluates to: " <> show (evalSym False model $ eval e1)
      putStrLn $ "rhs evaluates to: " <> show (evalSym False model $ eval e2)

-- | Build a synthesizer to find an equivalent expression for a given target
-- expression.
--
-- The synthesis problem can be formulated as finding values for holes in a
-- sketch expression that make it equivalent to a target expression.
-- Mathematically:
--
-- ∃s ∀x. target(x) = sketch(s,x)
--
-- Where:
--
-- * s represents the symbolic constants (holes) in the sketch that we want to
--   synthesize - these are the values we're trying to find
-- * x represents the free variables in both expressions - these are the inputs
--   that both expressions must handle equivalently
-- * target(x) is the target expression we want to match
-- * sketch(s,x) is the template expression with holes that we want to fill
--
-- We solve this using Counter-Example Guided Inductive Synthesis (CEGIS),
-- which:
--
-- 1. Proposes candidate values for holes
-- 2. Checks if the resulting expression matches the target for all inputs
-- 3. If not, uses counter-examples to refine the candidates
-- 4. Repeats until a solution is found or synthesis is proven impossible
--
-- For example, given target @a + a@ and sketch @a * c@:
--
-- * We want to synthesize a value for @c@ (the hole)
-- * Such that for any input value of @a@
-- * The equation @a + a = a * c@ holds
-- * CEGIS will discover @c = 2@ is the solution, since @a + a = a * 2@ for all
--   @a@
--
-- Grisette provides built-in CEGIS functionality to solve such synthesis
-- problems.
synthesisRewriteTarget :: (BasicSymPrim a) => Expr a -> UExpr a -> IO ()
synthesisRewriteTarget expr sketch = do
  r <- cegisForAll z3 expr $ cegisPostCond $ eval expr .== eval .# sketch
  case r of
    (_, CEGISSuccess model) -> do
      putStrLn $ "For the target expression: " <> show expr
      putStrLn "Successfully synthesized RHS:"
      print $ evalSym False model sketch
    (cex, failure) -> do
      putStrLn $ "Synthesis failed with error: " ++ show failure
      putStrLn $ "Counter example list: " ++ show cex

-- * Verification and synthesis examples

aPlusB :: IntExpr
aPlusB = Add (intVal "a") (intVal "b")

bPlusA :: IntExpr
bPlusA = Add (intVal "b") (intVal "a")

aPlusA :: IntExpr
aPlusA = Add (intVal "a") (intVal "a")

productOfSum :: IntExpr
productOfSum = Mul (intVal "a") (add (intVal "b") (intVal "c"))

sumOfProduct :: IntExpr
sumOfProduct =
  Add (mul (intVal "a") (intVal "b")) (mul (intVal "a") (intVal "c"))

allSum :: IntExpr
allSum = Add (intVal "a") (add (intVal "b") (intVal "c"))

xPlusX :: IntExpr
xPlusX = Add (intVal "x") (intVal "x")

xTimesC :: IntUExpr
xTimesC = mul (intVal "x") (intVal "c")

-- | We can use the `Fresh` monad to build sketches in a modular way.
--
-- Given a list of expressions, this function builds a new expression that
-- represents one of three symbolic choices:
--
-- * Addition of two expressions symbolically chosen from the input list
-- * Multiplication of two expressions symbolically chosen from the input list
-- * A single expression symbolically chosen from the input list
--
-- The `chooseUnionFresh` function generates fresh symbolic variables to
-- represent these choices. For example, given expressions [a, b, c], it
-- creates:
--
-- > {If cond@1 a {If cond@2 b c}}
--
-- Where cond@1 and cond@2 are unique symbolic boolean variables managed by the
-- `Fresh` monad. The solver can then determine the values of these boolean
-- variables to select the desired expressions. Each call to `chooseUnionFresh`
-- generates new variables, ensuring no overlap between different choices.
nextLevel :: [IntUExpr] -> Fresh IntUExpr
nextLevel exprs = do
  lhs <- chooseUnionFresh exprs
  rhs <- chooseUnionFresh exprs
  chooseUnionFresh [add lhs rhs, mul lhs rhs, lhs]

-- | Builds a hierarchical sketch for expression synthesis.
--
-- This function creates a two-level sketch where:
--
-- 1. The base level consists of atomic variables a, b, and c
-- 2. The middle level (l2, r2) applies nextLevel to combine these atoms
-- 3. The final level combines the middle level expressions
--
-- This creates a rich space of possible expressions by composing
-- additions, multiplications and variable selections at multiple levels.
getSketch :: Fresh IntUExpr
getSketch = do
  let atom = [intVal "a", intVal "b", intVal "c"]
  l2 <- nextLevel atom
  r2 <- nextLevel atom
  nextLevel [l2, r2]

sketch :: IntUExpr
sketch = runFresh getSketch "sketch"

main :: IO ()
main = do
  putStrLn "---- verifying a + b and b + a are equivalent ----"
  verifyEquivalent aPlusB bPlusA
  putStrLn "---- verifying a + a and a + b are not equivalent (should fail) ----"
  verifyEquivalent aPlusA aPlusB

  putStrLn "---- verifying productOfSum and sumOfProduct are equivalent ----"
  verifyEquivalent productOfSum sumOfProduct
  putStrLn "---- verifying productOfSum and allSum are equivalent (should fail) ----"
  verifyEquivalent productOfSum allSum

  putStrLn "---- synthesis x + x => x * 2 ----"
  synthesisRewriteTarget xPlusX xTimesC
  putStrLn "---- synthesis a * (b + c) => a * b + a * c ----"
  synthesisRewriteTarget productOfSum sketch