calc.hs

module Main(main) where

import Data.Maybe (fromJust)
import Data.List (nub)

import Text.Parsec
import Text.Parsec.Language
import qualified Text.Parsec.Token as P

-- For the SMT solver interface
import Data.SBV

-- integer arithmetic. division by 0 will be
-- defined as x/0 = 0. This is a "good" choice!
-- See here: https://hackage.haskell.org/package/sbv-5.12/docs/Data-SBV.html#g:28
data Op = Add | Sub | Mul | Div
       deriving Show

-- An equation is simply two expressions
-- we say are equal
data Eqn = Eqn E E deriving Show

-- data type for expressions; the novelty here is
-- that we allow arbitrary variables to occur freely
-- i.e.,         x+1*(12-y)+z
data E = N Integer
       | V String
       | BinOp Op E E
       | Paren E
       deriving Show

--------------------------------------------------------------------------------------
-- Parsing.. Mostly following the examples from
--   https://hackage.haskell.org/package/parsec-3.1.11/docs/Text-Parsec.html
-- not much creativity here, mostly boiler-plate
lexer   = P.makeTokenParser haskellDef
integer = P.integer    lexer
symbol  = P.symbol     lexer
parens  = P.parens     lexer
ident   = P.identifier lexer

eqn    = do e1 <- expr; symbol "="; e2 <- expr; return (Eqn e1 e2)
expr   = chainl1 term   plusMinus
term   = chainl1 factor timesDiv
factor = (Paren `fmap` parens expr) <|> (N `fmap` integer) <|> (V `fmap` ident)

plusMinus = (symbol "+" >> return (BinOp Add)) <|> (symbol "-" >> return (BinOp Sub))
timesDiv  = (symbol "*" >> return (BinOp Mul)) <|> (symbol "/" >> return (BinOp Div))

--------------------------------------------------------------------------------------
-- Now the fun begins. Let's create an SMT problem that corresponds to the equation
-- and ask the backend solver to "satisfy" it; i.e., give us a binding for variables
-- so that the equation is true.

-- collect the variables appearing in the equation
vars (Eqn a b) = nub $ go a ++ go b
  where go (N _)         = []
        go (V n)         = [n]
        go (BinOp _ a b) = go a ++ go b
        go (Paren e)     = go e

calc inp = sat $ do
  let equation = case parse eqn "" inp of
                   Left  s -> error (show s)
                   Right e -> e
      vs       = vars equation

  -- create a symbolic variable for each var, creating our "read-only" env
  env <- zip vs `fmap` mapM sInteger vs

  let -- translate each expression, using the env constructed above for variables
      translate (N i)         = literal i
      translate (V s)         = fromJust (lookup s env)
      translate (BinOp o a b) =  let a' = translate a
                                     b' = translate b
                                 in case o of
                                     Add -> a' + b'
                                     Sub -> a' - b'
                                     Mul -> a' * b'
                                     Div -> a' `sDiv` b'
      translate (Paren e)      = translate e

      -- compile each part separately, and assert equals
      walk (Eqn a b) = translate a .== translate b

  return $ walk equation

test inp = do putStrLn $ "== [" ++ inp ++ "]=================================="
              print =<< calc inp

main = do test "x=1+2"
          test "x=2+(-1)"
          test "x=(-1)*(-1)"
          test "x=(1+2)*3-4"
          test "x=1000+5"
          test "x=1+2*3-4"
          test "4=x+3*y-12"

Here's the output from ghci:

*Main> main
== [x=1+2]==================================
Satisfiable. Model:
  x = 3 :: Integer
== [x=2+(-1)]==================================
Satisfiable. Model:
  x = 1 :: Integer
== [x=(-1)*(-1)]==================================
Satisfiable. Model:
  x = 1 :: Integer
== [x=(1+2)*3-4]==================================
Satisfiable. Model:
  x = 5 :: Integer
== [x=1000+5]==================================
Satisfiable. Model:
  x = 1005 :: Integer
== [x=1+2*3-4]==================================
Satisfiable. Model:
  x = 3 :: Integer
== [4=x+3*y-12]==================================
Satisfiable. Model:
  x = 16 :: Integer
  y =  0 :: Integer