Experimenting with data-types

[LeventErkok]

This is hard to use and the solver doesn't really perform all that well; but here's one way to encode recursive tree like data-types in SBV and use the solver. Backend solver chokes for pretty much anything that's not really super easy. So, ROI is not really there. And the lack of proper support for data-types and pattern-matching makes this very painful.

{-# LANGUAGE DeriveAnyClass  #-}
{-# LANGUAGE TypeFamilies    #-}
{-# LANGUAGE TemplateHaskell #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Main(main) where

import Data.SBV
import Data.SBV.Either
import Data.SBV.Maybe
import Data.SBV.Control

import qualified Data.SBV.List as L

data E = ILit Integer
       | BLit Bool
       | Lt   E E
       | Add  E E
       | Mul  E E
       | If   E E E
       deriving Show

data Tag  = TILit | TBLit | TLt | TAdd | TMul | TIf
mkSymbolicEnumeration ''Tag

type Data = Either Integer Bool

type Elem  = Either Tag Data
data SE = SE (SList Elem)

-- Check that a given list corresponds to a valid E. Note that we don't do "type-checking." That is,
-- we only make sure the given list is the postfix equivalent of some freely generated element.
-- It would be possible to add "type-checking" to this as well, at least in this easy case, but
-- the complexity is hardly worth it.
validExpr :: SList Elem -> SBool
validExpr inp = go inp L.nil
  where go :: SList Elem -> SList (Either () Data) -> SBool
        go = smtFunction "validExpr" $ \tokens stk ->
                 let (tok, rest)  = L.uncons tokens

                     (s0,  stk1) = L.uncons stk
                     (s1,  stk2) = L.uncons stk1
                     (s2,  stk3) = L.uncons stk2

                     node :: SEither () Data
                     node = sLeft $ literal ()

                     isNode :: SEither () Data -> SBool
                     isNode = isLeft
              in
                 -- If tokens are exhausted, we must have one expression left on top of the stack
                 ite (L.null tokens) (L.length stk .== 1 .&& isLeft (stk L.!! 0))

                 -- Push the leaves in
               $ ite (isRight tok) (go rest (sRight (fromRight tok) L..: stk))

                 -- Analyze the constructors
               $ ite (tok .== sLeft sTILit .&& L.length stk .>= 1 .&& isRight s0 .&& isLeft  (fromRight s0))  (go rest (node L..: stk1))
               $ ite (tok .== sLeft sTBLit .&& L.length stk .>= 1 .&& isRight s0 .&& isRight (fromRight s0))  (go rest (node L..: stk1))
               $ ite (tok .== sLeft sTLt   .&& L.length stk .>= 2 .&& isNode  s0 .&& isNode s1 )              (go rest (node L..: stk2))
               $ ite (tok .== sLeft sTAdd  .&& L.length stk .>= 2 .&& isNode  s0 .&& isNode s1 )              (go rest (node L..: stk2))
               $ ite (tok .== sLeft sTMul  .&& L.length stk .>= 2 .&& isNode  s0 .&& isNode s1 )              (go rest (node L..: stk2))
               $ ite (tok .== sLeft sTIf   .&& L.length stk .>= 3 .&& isNode  s0 .&& isNode s1 .&& isNode s2) (go rest (node L..: stk3))
                 sFalse

instance Queriable IO SE where
  type QueryResult SE = E

  create = do r@(SE l) <- SE <$> freshVar_
              constrain $ validExpr l
              pure r

  project (SE x) = se2e <$> getValue x
    where se2e :: [Either Tag Data] -> E
          se2e inp = go inp []
            where go :: [Either Tag Data] -> [Either E Data] -> E

                  -- If tokens are exhausted, we must have one expression left on top of the stack
                  go [] [Left tos] = tos

                  -- Pushing leaves in
                  go (Right d : rest) stk = go rest (Right d : stk)

                  -- Popping constructors out
                  go (Left TILit : rest) (                    Right (Left  i) : stk)  = go rest (Left (ILit i)        : stk)
                  go (Left TBLit : rest) (                    Right (Right b) : stk)  = go rest (Left (BLit b)        : stk)
                  go (Left TLt   : rest) (          Left a2 : Left         a1 : stk)  = go rest (Left (Lt   a1 a2)    : stk)
                  go (Left TAdd  : rest) (          Left a2 : Left         a1 : stk)  = go rest (Left (Add  a1 a2)    : stk)
                  go (Left TMul  : rest) (          Left a2 : Left         a1 : stk)  = go rest (Left (Mul  a1 a2)    : stk)
                  go (Left TIf   : rest) (Left a3 : Left a2 : Left         a1 : stk)  = go rest (Left (If   a1 a2 a3) : stk)

                  -- Shouldn't happen, unless we screw something up
                  go xs stk = error $ unlines [ ""
                                              , "*** SE2E, failed."
                                              , "*   Input      : " ++ show inp
                                              , "*   Tokens left: " ++ show xs
                                              , "*   Stack      : " ++ show stk
                                              ]


  embed e = pure $ SE $ literal (e2se e)
    where e2se :: E -> [Either Tag Data]
          e2se (ILit i)     = [Right (Left  i), Left TILit]
          e2se (BLit b)     = [Right (Right b), Left TBLit]
          e2se (Lt   a b)   = e2se a ++ e2se b           ++ [Left TLt]
          e2se (Add  a b)   = e2se a ++ e2se b           ++ [Left TAdd]
          e2se (Mul  a b)   = e2se a ++ e2se b           ++ [Left TMul]
          e2se (If   a b c) = e2se a ++ e2se b ++ e2se c ++ [Left TIf]

eval :: SE -> SMaybe (Either Integer Bool)
eval (SE inp) = go inp L.nil
  where go :: SList Elem -> SList (Either Integer Bool) -> SMaybe (Either Integer Bool)
        go = smtFunction "evalExpr" $ \tokens stk ->
                 let (tok, rest)  = L.uncons tokens

                     (s0,  stk1) = L.uncons stk
                     (s1,  stk2) = L.uncons stk1
                     (s2,  stk3) = L.uncons stk2

              in
                 -- If tokens are exhausted, we must have one expression left on top of the stack. return that.
                 ite (L.null tokens .&& L.length stk .== 1) (sJust s0)

                 -- Push the leaves in
               $ ite (isRight tok) (go rest (fromRight tok L..: stk))

                 -- Analyze the constructors
               $ ite (tok .== sLeft sTILit .&& L.length stk .>= 1 .&& isLeft  s0)                               (go rest (s0                                                        L..: stk1))
               $ ite (tok .== sLeft sTBLit .&& L.length stk .>= 1 .&& isRight s0)                               (go rest (s0                                                        L..: stk1))
               $ ite (tok .== sLeft sTLt   .&& L.length stk .>= 2 .&& isLeft  s0 .&& isLeft  s1)                (go rest (sRight (fromLeft s1 .< fromLeft s0)                       L..: stk2))
               $ ite (tok .== sLeft sTAdd  .&& L.length stk .>= 2 .&& isLeft  s0 .&& isLeft  s1)                (go rest (sLeft  (fromLeft s1 +  fromLeft s0)                       L..: stk2))
               $ ite (tok .== sLeft sTMul  .&& L.length stk .>= 2 .&& isLeft  s0 .&& isLeft  s1)                (go rest (sLeft  (fromLeft s1 *  fromLeft s0)                       L..: stk2))
               $ ite (tok .== sLeft sTIf   .&& L.length stk .>= 3 .&& isLeft  s0 .&& isLeft  s1 .&& isRight s2) (go rest (sLeft  (ite (fromRight s2) (fromLeft  s1) (fromLeft  s0)) L..: stk3))
               $ ite (tok .== sLeft sTIf   .&& L.length stk .>= 3 .&& isRight s0 .&& isRight s1 .&& isRight s2) (go rest (sRight (ite (fromRight s2) (fromRight s1) (fromRight s0)) L..: stk3))
                 sNothing

ex1 :: Bool -> IO ()
ex1 v = print =<< runSMTWith z3{verbose=v} p
 where p = do e@(SE l) <- SE <$> free_
              constrain $ validExpr l
              constrain $ eval e .== sJust (sLeft 12)
              query $ do ensureSat
                         project e

-- A much simpler example that goes through, which does a simple eval
ex2 :: Bool -> IO ()
ex2 v = print =<< runSMTWith z3{verbose=v} p
 where p = do e@(SE l) <- SE <$> free_
              constrain $ validExpr l
              constrain $ eval e .== eval e

              res <- free_

              query $ do expr <- embed $ If (Lt (ILit 13) (ILit 5)) (ILit 12) (ILit 21)

                         constrain $ res .== eval expr
                         ensureSat
                         getValue res

main :: IO ()
main = do let verbose = False
          -- ex1 verbose  -- Too difficult
          ex2 verbose


_unused :: STag
_unused = error "Unused" ex1