# Language GADTs , ViewPatterns , RankNTypes #
# LANGUAGE GeneralizedNewtypeDeriving , DeriveFunctor #

  Parser combinators for reading back a model from the solver

module Talos.SymExec.ModelParser
  ( ModelP, evalModelP
  , pByte, pNumber, pBytes, pMUnit, pIndexed, pSeq, pValue
  ) where

import           Control.Applicative
import           Control.Monad       (ap, guard)
import           Data.ByteString     (ByteString)
import qualified Data.ByteString     as BS
import           Data.Char           (isDigit)
import           Data.List           (foldl')
import           Data.Map            (Map)
import qualified Data.Map            as Map
import qualified Data.Vector         as V
import           Data.Word
import           SimpleSMT           (SExpr (..))
import           Text.Read           (readMaybe)

import           Daedalus.Core
import           Daedalus.PP         (showPP)
import           Daedalus.Panic
import qualified Daedalus.Value      as I

--------------------------------------------------------------------------------
  Model Combinators
--
-- Only what is needed by the relation type generated in SymExec.hs
-- (sum, product, bytes, unit) 

pBytes :: ModelP ByteString
pBytes = pSExpr $ do
  pExact "bytes"
  pByteString

pMUnit :: ModelP ()
pMUnit = pExact "munit"

pIndexed :: (Int -> ModelP a) -> ModelP a
pIndexed rest = pSExpr $ do
  pExact "indexed"
  n <- pNumber
  rest (fromIntegral n)

   : : ( Int - > ModelP a ) - > ModelP a
-- pSeq m = pSExpr $ do
--   pExact "seq"
--   n  <- pNumber -- count
--   pArray n (m (fromIntegral n))

pSeq :: ModelP a -> ModelP b -> ModelP (a, b)
pSeq pl pr = pSExpr $ do
  pExact "seq"
  (,) <$> pl <*> pr

-- -----------------------------------------------------------------------------
  Monadic parser

-- This is complicated by the fact that the solver returns
-- non-normalised values, including references to other variable, etc.

type VEnv = Map String SExpr

-- This is a bit hacky
data ModelPState =
  ModelPState {  vEnv   :: VEnv
               , sexps  :: [SExpr]
               }

newtype ModelP a = ModelP { runModelP :: ModelPState -> [(a, ModelPState)] }
  deriving (Functor)

instance Applicative ModelP where
  pure v   = ModelP $ \s -> [(v, s)]
  (<*>)    = ap

instance Monad ModelP where
  list 
    (a, s') <- runModelP m s
    runModelP (f a) s'
    -- -- Not exactly pretty, but it suffices for out use case
    -- concat <$> mapM (\(a, s') -> runModelP (f a) s') ress

instance Alternative ModelP where
  empty    = ModelP $ const []
  m <|> m' = ModelP $ \s -> runModelP m s ++ runModelP m' s

-- instance MonadIO ModelP where
    liftIO m = ModelP $ \s - > m > > = \a - > pure [ ( a , s ) ]

evalModelP' :: ModelP a -> ModelPState -> [a]
evalModelP' m s = map fst (runModelP m s)

evalModelP :: ModelP a -> SExpr -> [a]
evalModelP m sexpr = evalModelP' m (ModelPState Map.empty [sexpr])

--------------------------------------------------------------------------------
-- Primitive combinators

pRawHead :: ModelP SExpr
pRawHead = ModelP $ \st -> do
  case sexps st of
    []     -> []
    s : ss -> [(s, st { sexps = ss })]

  Runs the parser p on the first element , taking care of let - related
-- and 'as' annoyances
pHead :: (String -> ModelP a) -> -- ^ This is the atom case, it is in ModelP as it can fail.
         ModelP a -> -- ^ THe list case 
         ModelP a
pHead pA pS = do
  s <- pRawHead
  ModelP $ \st -> do -- list monad
    r <- handleLet st s
    pure (r, st)
  where
    handleLet st (List [ident -> Just "let", List idents, b]) =
      handleLet (st { vEnv = foldl' addBind (vEnv st) idents }) b
    handleLet st (List [ident -> Just "as", s, _]) = handleLet st s
    -- FIXME: This does a map lookup for each atom in the sexp
    handleLet st (Atom a)
      | Just s <- Map.lookup a (vEnv st) = handleLet st s
      | otherwise = evalModelP' (pA a) (st { sexps = [] })
    handleLet st (List ss') = evalModelP' pS (st { sexps = ss' })

    addBind env (List [ident -> Just v, e]) = Map.insert v e env
    addBind _    v = panic "handleLet unexpected let binding" [show v]

pAtom :: ModelP String
pAtom = pHead pure empty

pExact :: String -> ModelP ()
pExact a = pAtom >>= guard . (==) a

pSExpr :: ModelP a -> ModelP a
pSExpr = pHead (const empty)

withSExprs :: [SExpr] -> ModelP a -> ModelP a
withSExprs ss m = ModelP $ \st -> do
  r <- evalModelP' m (st { sexps = ss })
  pure (r, st)

pArray :: Integer -> ModelP a -> ModelP a
pArray count p = do
  (els, arr) <- go
  let v = V.replicate (fromInteger count) arr V.// els
  withSExprs (V.toList v) p
  where
    -- FIXME: maybe use a mutable vector here?
    base :: ModelP ([(Int, SExpr)], SExpr)
    base = pSExpr $ do
      pExact "const"
      s <- pRawHead
      pure ([], s)

    go :: ModelP ([(Int, SExpr)], SExpr)
    go =
      base <|>
      (pSExpr $
        do pExact "store"
           (els, arr) <- go
           i   <- pNumber
           s   <- pRawHead
           pure (if i < count then (fromInteger i, s) : els else els, arr))

-- FIXME: do we need this?

   : : String - > ModelP SExpr
   a = do
     < - ModelP $ asks ( Map.lookup a . vEnv )
    case  of
--     Just v -> pure v
--     Nothing -> do
--       s <- ModelP $ asks solver
--       res <- liftIO $ command s $ fun "get-value" [List [S.const a]]
--       case res of
--         List [List [_, v]] -> pure v
          _ - > panic ( unlines
                  [ " Unexpected response from the SMT solver : "
                  , "   Exptected : a value "
--                 , "  Result: " ++ showsSExpr res ""
--                 ]) []


--------------------------------------------------------------------------------
-- Combinators
--
-- Only what is needed by the relation type generated in SymExec.hs
-- (sum, product, bytes, unit) 

  pTuple : : ModelP a - > ModelP b - > ModelP ( a , b )
  pTuple l r = pSExpr $ do
--   pExact "mk-tuple"
--   (,) <$> l <*> r

  pUnit : : ModelP ( )
  pUnit = pExact " unit "

-- pSum :: ModelP a -> ModelP a -> ModelP a
-- pSum l r = pSExpr $ do
--   (pExact "inl" *> l) <|> (pExact "inr" *> r)

pNumber :: ModelP Integer
pNumber = pAtom >>= go
  where
    go ('#' : 'x' : rest) = do
      let str = "0x" ++ rest
      case readMaybe str of
        Nothing -> error ("Couldn't read " ++ show str)
        Just v  -> pure v
    go ('#' : 'b' : rest) = do
      case binLit rest of
        Nothing -> error ("Couldn't read " ++ show rest)
        Just v  -> pure v
    go ds | all isDigit ds = do
      case readMaybe ds of
        Nothing -> error ("Couldn't read " ++ show ds)
        Just v  -> pure v
    go _ = empty

    -- copied from simple-smt
    binLit cs = do ds <- mapM binDigit cs
                   return $ sum $ zipWith (*) (reverse ds) powers2
    powers2   = 1 : map (2 *) powers2
    binDigit '0' = Just 0
    binDigit '1' = Just 1
    binDigit _   = Nothing

pByte :: ModelP Word8
pByte = fromIntegral <$> pNumber

  This is nicer , we could use  directly
pList :: ModelP a -> ModelP [a]
pList p = ([] <$ pExact "nil") <|>
          (pSExpr $ do pExact "insert"
                       (:) <$> p <*> pList p)

pByteString :: ModelP ByteString
pByteString = BS.pack <$> pList pByte

pBool :: ModelP Bool
pBool = pAtom >>= go
  where
    go "true"  = pure True
    go "false" = pure False
    go _       = empty

-- -----------------------------------------------------------------------------
-- Values

pValue :: Type -> ModelP I.Value
pValue ty0 = go ty0
  where
  c.f . symExecTy
    go ty =
      case ty of
        TStream -> unimplemented
        TUInt (TSize n) -> I.vUInt (fromIntegral n) <$> pNumber
        TUInt _ -> unimplemented
        TSInt (TSize n) -> I.vSInt' (fromIntegral n) <$> pNumber
        TSInt _ -> unimplemented
  FIXME
        TBool -> I.VBool <$> pBool
        TUnit -> pure I.vUnit
        TArray _ -> unimplemented
        TMaybe _ty' -> unimplemented
        TMap {} -> unimplemented
        TBuilder {} -> unimplemented
        TIterator {} -> unimplemented
        TUser _ut -> unimplemented
        TParam {} -> unimplemented

    unimplemented = panic "sexprToValue: Unimplemented" [ showPP ty0 ]

--------------------------------------------------------------------------------
-- Helpers

ident :: SExpr -> Maybe String
ident (Atom a) = Just a
ident (List [ident -> Just "as", v, _]) = ident v
ident _ = Nothing

-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------
-- --------------------------------------------------------------------------------



   : : Solver - > VEnv - > ModelP a - > String - > IO a
  getValue s e  ( ( Map . ! ? ) e - > Just r ) = runModelP    
  getValue s e  v = do
--   res <- command s $ fun "get-value" [List [fun v []]]
--   case res of
--     List [List [_, v]] -> runModelP pOne s e v
      _ - > fail $ unlines
                  [ " Unexpected response from the SMT solver : "
                  , "   Exptected : a value "
--                 , "  Result: " ++ showsSExpr res ""
--                 ]

  pBool : : ModelP 
  pBool = pLisp ( ModelP go )
--   where
--     go s env (ident -> Just "true") = pure True
--     go s env (ident -> Just "false") = pure False
      go s env ( ident - > Just v ) =  s env  v
      go s env e = error "   

-- pNumber :: ModelP Integer
-- pNumber = pLisp (ModelP go)
--   where
--     go s env (ident -> Just ('#' : 'x' : rest)) = do
        let str = " 0x " + + rest
        case   of
--         Nothing -> error ("Couldn't read " ++ show str)
--         Just v  -> pure v
--     go s env (ident -> Just ('#' : 'b' : rest)) = do
--       case binLit rest of
--         Nothing -> error ("Couldn't read " ++ show rest)
--         Just v  -> pure v
--     go s env (ident -> Just ds) | all isDigit ds = do
        case   of
--         Nothing -> error ("Couldn't read " ++ show ds)
--         Just v  -> pure v
      go s env ( ident - > Just v ) =  s env pNumber v
      go s env e = error " Unparseable byte "

--     -- copied from simple-smt
--     binLit cs = do ds <- mapM binDigit cs
--                    return $ sum $ zipWith (*) (reverse ds) powers2
      powers2    = 1 : map ( 2 * ) powers2
--     binDigit '0' = Just 0
      binDigit ' 1 ' = Just 1
--     binDigit _   = Nothing

   : : ModelP Word8
-- pByte = fromIntegral <$> pNumber

  pList : : ModelP a - > ModelP [ a ]
-- pList p = pLisp (ModelP go)
--   where
--     go s env e = 
--       case e of
--         (ident -> Just "nil") -> pure []
          -- This would be nicer in pLisp , but we ca n't distinguish
--         -- between new model variables and regular consts.  Maybe
--         -- check for 'foo!n' symbols?
          ( ident - > Just v )      - >  s env ( pList p ) v
--         List [ident -> Just "insert", e', more] ->
--           (:) <$> runModelP p s env e' <*> go s env more

--         other  -> error $ "Not a list" ++ show e

-- pListWithLength :: ModelP a -> ModelP [a]
  pListWithLength p = pLisp ( ModelP go )
--   where
--     go s env e =
--       case e of
          ( ident - > Just v )      - >  s env ( pListWithLength p ) v
--         List [ident -> Just "mk-ListWithLength", l, _] -> runModelP (pList p) s env l
--         other  -> error $ "Not a list with length " ++ show e



  getBytes : : Solver - > SExpr - > IO ByteString
  getBytes s e = BS.pack < $ > runModelP ( pListWithLength  ) s Map.empty e



-- -- -----------------------------------------------------------------------------
-- -- Top level

  -- callG : : Solver - > Integer - > Name - > Type - > [ Type ] - > [ I.Value ] - > IO ByteString
-- -- callG s offset name resT argTs args = inNewScope s $ do
  --    b0 < - declare s " |bytes before| " (   )
-- --   o0 <- declare s "|oracle before|" tOracle
-- --   let i0  = S.fun "mk-input" [b0, S.int offset]
-- --       st0 = S.fun "mk-state" [i0, o0]
  --         s " res " ( symExecTy resT )
-- --   st' <- declare s "|state after|" tState

-- --   let args' = zipWith symExecValue argTs args
-- --       call = S.fun (ruleName name) (args' ++ [st0])

-- --   -- We don't really care about index, only that it is increasing
-- --   -- assert s (eq (S.fun "index" [o0]) (S.tInt 0)

  --    assert s ( eq ( S.fun " bytes " [ S.fun " input " [ st ' ] ] ) (  ( S.const " nil " ) (   ) ) )

-- --   -- assert that calling the parser doesn't fail (we don't care about
-- --   -- the resulting input (maybe?)
-- --   -- assert s (S.not (fun "is-fail" [call]))
  --    assert s ( eq call ( fun " success " [ st ' ,  ] ) )
-- --   r <- check s
-- --   case r of
-- --     Sat -> pure ()
-- --     -- FIXME: pretty gross
  --          - > error "  "
-- --     Unknown -> error "Unknown"

-- --   -- v <- S.getExpr s rv
-- --   -- print v

-- --   getByteModel s b0


-- --------------------------------------------------------------------------------
-- -- Constraint satisfaction

-- -- solvePathConstraints :: Solver
  --                       - > Map ( TCName K.Value ) I.Value
  --                       - > Map ( TCName K.Value ) VarHandle
  --                       - > Set VarHandle
  --                       - > PathConstraints
-- --                      -> IO PathConstraintSolution
  -- solvePathConstraints s   bsVs pcs = inNewScope s $ do
-- --   -- print pcs
-- --   -- declare all the vars
  --    mapM _ declVar ( Map.toList  )
  --    let env = Env { symVars = ( symExecVarHandle < $ >  ) ` Map.union ` env0 }
-- --   mapM_ declBVar (Set.toList bsVs)
-- --   -- print (show (map (symExecPathConstraint env) pcs))
  --    mapM _ ( S.assert s . symExecPathConstraint env ) pcs
-- --   r <- check s
-- --   case r of
-- --     Sat -> pure ()
-- --     -- FIXME: pretty gross
  --          - > error "  "
-- --     Unknown -> error "Unknown"

  --    rvMap < - Map.fromList < $ > mapM  ' ( Map.toList  )
-- --   bsMap <- sequence (Map.fromSet (getBytes s . symExecVarHandle) bsVs)

-- --   pure (rvMap, bsMap) -- FIXME
-- --   where
  --      env0                = Map.mapWithKey ( \tcn - > symExecValue ( tcType tcn ) ) vEnv
  --      declVar ( tcn ,  ) = S.declare s ( varHandleName  )   ( symExecTy ( tcType tcn ) )
  --      declBVar bsId       =  s ( varHandleName bsId ) ( tListWithLength  )

  --       ' ( tcn ,  ) = do
-- --       v <- runModelP (sexprToValue (tcType tcn)) s Map.empty (symExecVarHandle vId)
-- --       -- print $ pp tcn <+> pp v
  --        pure (  , v )

-- -- Turns a path constraint into a predicate.
  -- symExecPathConstraint : : Env - > PathConstraint - > SExpr
  -- symExecPathConstraint e (  bsId vId tc ) =
-- --   case texprValue tc of
  --      -- These are the ' simple ' terms identified by  : all other terms are handled elsewhere
-- --     TCMatchBytes _ws v ->
-- --       S.and (S.eq (symExecVarHandle bsId) (symExecVarHandle vId))
-- --             (S.eq (symExecVarHandle vId)  (symExecV v))

  --       v         - >
  --        S.and ( S.eq ( symExecVarHandle bsId ) (   ) )
-- --             (S.eq (symExecVarHandle vId)  (symExecV v))

  --      TCGetByte { }     - > getByte
  --             - >
-- --       S.and getByte
-- --             (symExecP p (symExecVarHandle vId))

-- --     _               -> error ("BUG: saw a non-simple term: " ++ show (pp tc))
-- --     where
  --        getByte = S.eq ( symExecVarHandle bsId ) ( sCons ( symExecVarHandle vId ) (   ) )

  -- symExecPathConstraint e ( AssertionConstraint assn ) =
-- --   case assn of
  --          symExecV g

  symExecValue : : Type - > I.Value - >   SExpr
  symExecValue   =
    case  of
      VUInt n v   - > S.value ( S.Bits n v )
      VSInt n v   - > S.value ( S.Bits n v )
--     VInteger v -> S.value (S.Int    v)
      VBool b     - > S.value ( S.Bool b )
--     VUnionElem {} -> unimplemented
--     VStruct    {} -> unimplemented
--     VArray     vs | Just t <- isArray typ ->
        mkList ( map ( symExecValue t ) ( Vector.toList vs ) )
--     VArray     {} -> error "BUG"
--     VMaybe v | Just t <- isMaybe typ ->
--       case v of
--         Nothing -> sNothing (symExecTy t)
--         Just v  -> sJust (symExecValue t v)
--     VMaybe _   -> error "BUG"
       { }     - > unimplemented
--     VStream {} -> unimplemented
--   where
      unimplemented = error $ " symExecValue : Unimplemented " + + show 
      mkList : : [ SExpr ] - > SExpr
       = symExecTy 
--     mkList xs = fun "mk-ListWithLength"
                  [ foldr ( \a b - > fun " insert " [ a , b ] ) (  ( S.const " nil " ) ( tList  ) ) xs
                  , S.int ( fromIntegral $ length xs )
--                 ]

-- -- Constructs a parser for the value of the given type
-- sexprToValue :: Type -> ModelP Value
-- sexprToValue = go Map.empty
--   where
      -- c.f . symExecTy
--     go env ty =
--       case ty of
          TVar x | Just t < - env Map . ! ? x - > go env t
          TVar x - > error ( " Unknown type variable " + + show ( pp x ) )

--         TCon tname args -> unimplemented -- S.fun (symExecTCTyName tname) (map symExecTy args)
--         Type t ->
--           case t of
               { } - > error " Should n't happen ( symExecTy :  ) "
              TStream      - > error " Unimplemented "
                  error " Unimplemented " -- Should n't happen ?
              TNum { }      - > error " Should n't happen ( symExecTy : TNum ) "
--             TUInt (Type (TNum n)) -> VUInt (fromIntegral n) <$> pNumber
--             TUInt {}    -> error "Shouldn't happen (symExecTy: TUInt)"
              TSInt ( Type ( TNum n ) ) - > VSInt ( fromIntegral n ) < $ > pNumber
--             TSInt {}    -> error "Shouldn't happen (symExecTy: TSInt)"
--             TInteger    -> VInteger <$> pNumber
                - > VBool < $ > pBool
--             TUnit       -> unimplemented
              TArray t '    - > VArray . Vector.fromList < $ > pListWithLength ( sexprToValue t ' )
              TMaybe t '    - > unimplemented -- VMaybe < $ > pMaybe ( sexprToValue env t ' )
              TMap kt     unimplemented -- tMap ( symExecTy ' env kt ) ( symExecTy ' env vt )

--     unimplemented = error $ "sexprToValue: Unimplemented" 

-- -- -----------------------------------------------------------------------------
-- -- Probably belongs elsewhere


------------------------------------------------------------------------------

 Only what is needed by the relation type generated in SymExec.hs
 (sum, product, bytes, unit) 
 pSeq m = pSExpr $ do
   pExact "seq"
   n  <- pNumber -- count
   pArray n (m (fromIntegral n))
 -----------------------------------------------------------------------------
 This is complicated by the fact that the solver returns
 non-normalised values, including references to other variable, etc.
 This is a bit hacky
 -- Not exactly pretty, but it suffices for out use case
 concat <$> mapM (\(a, s') -> runModelP (f a) s') ress
 instance MonadIO ModelP where
------------------------------------------------------------------------------
 Primitive combinators
 and 'as' annoyances
 ^ This is the atom case, it is in ModelP as it can fail.
 ^ THe list case 
 list monad
 FIXME: This does a map lookup for each atom in the sexp
 FIXME: maybe use a mutable vector here?
 FIXME: do we need this?
     Just v -> pure v
     Nothing -> do
       s <- ModelP $ asks solver
       res <- liftIO $ command s $ fun "get-value" [List [S.const a]]
       case res of
         List [List [_, v]] -> pure v
                 , "  Result: " ++ showsSExpr res ""
                 ]) []
------------------------------------------------------------------------------
 Combinators

 Only what is needed by the relation type generated in SymExec.hs
 (sum, product, bytes, unit) 
   pExact "mk-tuple"
   (,) <$> l <*> r
 pSum :: ModelP a -> ModelP a -> ModelP a
 pSum l r = pSExpr $ do
   (pExact "inl" *> l) <|> (pExact "inr" *> r)
 copied from simple-smt
 -----------------------------------------------------------------------------
 Values
------------------------------------------------------------------------------
 Helpers
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
 --------------------------------------------------------------------------------
   res <- command s $ fun "get-value" [List [fun v []]]
   case res of
     List [List [_, v]] -> runModelP pOne s e v
                 , "  Result: " ++ showsSExpr res ""
                 ]
   where
     go s env (ident -> Just "true") = pure True
     go s env (ident -> Just "false") = pure False
 pNumber :: ModelP Integer
 pNumber = pLisp (ModelP go)
   where
     go s env (ident -> Just ('#' : 'x' : rest)) = do
         Nothing -> error ("Couldn't read " ++ show str)
         Just v  -> pure v
     go s env (ident -> Just ('#' : 'b' : rest)) = do
       case binLit rest of
         Nothing -> error ("Couldn't read " ++ show rest)
         Just v  -> pure v
     go s env (ident -> Just ds) | all isDigit ds = do
         Nothing -> error ("Couldn't read " ++ show ds)
         Just v  -> pure v
     -- copied from simple-smt
     binLit cs = do ds <- mapM binDigit cs
                    return $ sum $ zipWith (*) (reverse ds) powers2
     binDigit '0' = Just 0
     binDigit _   = Nothing
 pByte = fromIntegral <$> pNumber
 pList p = pLisp (ModelP go)
   where
     go s env e = 
       case e of
         (ident -> Just "nil") -> pure []
 This would be nicer in pLisp , but we ca n't distinguish
         -- between new model variables and regular consts.  Maybe
         -- check for 'foo!n' symbols?
         List [ident -> Just "insert", e', more] ->
           (:) <$> runModelP p s env e' <*> go s env more
         other  -> error $ "Not a list" ++ show e
 pListWithLength :: ModelP a -> ModelP [a]
   where
     go s env e =
       case e of
         List [ident -> Just "mk-ListWithLength", l, _] -> runModelP (pList p) s env l
         other  -> error $ "Not a list with length " ++ show e
 -- -----------------------------------------------------------------------------
 -- Top level
 callG : : Solver - > Integer - > Name - > Type - > [ Type ] - > [ I.Value ] - > IO ByteString
 -- callG s offset name resT argTs args = inNewScope s $ do
    b0 < - declare s " |bytes before| " (   )
 --   o0 <- declare s "|oracle before|" tOracle
 --   let i0  = S.fun "mk-input" [b0, S.int offset]
 --       st0 = S.fun "mk-state" [i0, o0]
         s " res " ( symExecTy resT )
 --   st' <- declare s "|state after|" tState
 --   let args' = zipWith symExecValue argTs args
 --       call = S.fun (ruleName name) (args' ++ [st0])
 --   -- We don't really care about index, only that it is increasing
 --   -- assert s (eq (S.fun "index" [o0]) (S.tInt 0)
    assert s ( eq ( S.fun " bytes " [ S.fun " input " [ st ' ] ] ) (  ( S.const " nil " ) (   ) ) )
 --   -- assert that calling the parser doesn't fail (we don't care about
 --   -- the resulting input (maybe?)
 --   -- assert s (S.not (fun "is-fail" [call]))
    assert s ( eq call ( fun " success " [ st ' ,  ] ) )
 --   r <- check s
 --   case r of
 --     Sat -> pure ()
 --     -- FIXME: pretty gross
          - > error "  "
 --     Unknown -> error "Unknown"
 --   -- v <- S.getExpr s rv
 --   -- print v
 --   getByteModel s b0
 --------------------------------------------------------------------------------
 -- Constraint satisfaction
 -- solvePathConstraints :: Solver
                       - > Map ( TCName K.Value ) I.Value
                       - > Map ( TCName K.Value ) VarHandle
                       - > Set VarHandle
                       - > PathConstraints
 --                      -> IO PathConstraintSolution
 solvePathConstraints s   bsVs pcs = inNewScope s $ do
 --   -- print pcs
 --   -- declare all the vars
    mapM _ declVar ( Map.toList  )
    let env = Env { symVars = ( symExecVarHandle < $ >  ) ` Map.union ` env0 }
 --   mapM_ declBVar (Set.toList bsVs)
 --   -- print (show (map (symExecPathConstraint env) pcs))
    mapM _ ( S.assert s . symExecPathConstraint env ) pcs
 --   r <- check s
 --   case r of
 --     Sat -> pure ()
 --     -- FIXME: pretty gross
          - > error "  "
 --     Unknown -> error "Unknown"
    rvMap < - Map.fromList < $ > mapM  ' ( Map.toList  )
 --   bsMap <- sequence (Map.fromSet (getBytes s . symExecVarHandle) bsVs)
 --   pure (rvMap, bsMap) -- FIXME
 --   where
      env0                = Map.mapWithKey ( \tcn - > symExecValue ( tcType tcn ) ) vEnv
      declVar ( tcn ,  ) = S.declare s ( varHandleName  )   ( symExecTy ( tcType tcn ) )
      declBVar bsId       =  s ( varHandleName bsId ) ( tListWithLength  )
       ' ( tcn ,  ) = do
 --       v <- runModelP (sexprToValue (tcType tcn)) s Map.empty (symExecVarHandle vId)
 --       -- print $ pp tcn <+> pp v
        pure (  , v )
 -- Turns a path constraint into a predicate.
 symExecPathConstraint : : Env - > PathConstraint - > SExpr
 symExecPathConstraint e (  bsId vId tc ) =
 --   case texprValue tc of
      -- These are the ' simple ' terms identified by  : all other terms are handled elsewhere
 --     TCMatchBytes _ws v ->
 --       S.and (S.eq (symExecVarHandle bsId) (symExecVarHandle vId))
 --             (S.eq (symExecVarHandle vId)  (symExecV v))
       v         - >
        S.and ( S.eq ( symExecVarHandle bsId ) (   ) )
 --             (S.eq (symExecVarHandle vId)  (symExecV v))
      TCGetByte { }     - > getByte
             - >
 --       S.and getByte
 --             (symExecP p (symExecVarHandle vId))
 --     _               -> error ("BUG: saw a non-simple term: " ++ show (pp tc))
 --     where
        getByte = S.eq ( symExecVarHandle bsId ) ( sCons ( symExecVarHandle vId ) (   ) )
 symExecPathConstraint e ( AssertionConstraint assn ) =
 --   case assn of
          symExecV g
     VInteger v -> S.value (S.Int    v)
     VUnionElem {} -> unimplemented
     VStruct    {} -> unimplemented
     VArray     vs | Just t <- isArray typ ->
     VArray     {} -> error "BUG"
     VMaybe v | Just t <- isMaybe typ ->
       case v of
         Nothing -> sNothing (symExecTy t)
         Just v  -> sJust (symExecValue t v)
     VMaybe _   -> error "BUG"
     VStream {} -> unimplemented
   where
     mkList xs = fun "mk-ListWithLength"
                 ]
 -- Constructs a parser for the value of the given type
 sexprToValue :: Type -> ModelP Value
 sexprToValue = go Map.empty
   where
 c.f . symExecTy
     go env ty =
       case ty of
         TCon tname args -> unimplemented -- S.fun (symExecTCTyName tname) (map symExecTy args)
         Type t ->
           case t of
 Should n't happen ?
             TUInt (Type (TNum n)) -> VUInt (fromIntegral n) <$> pNumber
             TUInt {}    -> error "Shouldn't happen (symExecTy: TUInt)"
             TSInt {}    -> error "Shouldn't happen (symExecTy: TSInt)"
             TInteger    -> VInteger <$> pNumber
             TUnit       -> unimplemented
 VMaybe < $ > pMaybe ( sexprToValue env t ' )
 tMap ( symExecTy ' env kt ) ( symExecTy ' env vt )
     unimplemented = error $ "sexprToValue: Unimplemented" 
 -- -----------------------------------------------------------------------------
 -- Probably belongs elsewhere
# Language GADTs , ViewPatterns , RankNTypes #
# LANGUAGE GeneralizedNewtypeDeriving , DeriveFunctor #

  Parser combinators for reading back a model from the solver

module Talos.SymExec.ModelParser
  ( ModelP, evalModelP
  , pByte, pNumber, pBytes, pMUnit, pIndexed, pSeq, pValue
  ) where

import           Control.Applicative
import           Control.Monad       (ap, guard)
import           Data.ByteString     (ByteString)
import qualified Data.ByteString     as BS
import           Data.Char           (isDigit)
import           Data.List           (foldl')
import           Data.Map            (Map)
import qualified Data.Map            as Map
import qualified Data.Vector         as V
import           Data.Word
import           SimpleSMT           (SExpr (..))
import           Text.Read           (readMaybe)

import           Daedalus.Core
import           Daedalus.PP         (showPP)
import           Daedalus.Panic
import qualified Daedalus.Value      as I

  Model Combinators

pBytes :: ModelP ByteString
pBytes = pSExpr $ do
  pExact "bytes"
  pByteString

pMUnit :: ModelP ()
pMUnit = pExact "munit"

pIndexed :: (Int -> ModelP a) -> ModelP a
pIndexed rest = pSExpr $ do
  pExact "indexed"
  n <- pNumber
  rest (fromIntegral n)

   : : ( Int - > ModelP a ) - > ModelP a

pSeq :: ModelP a -> ModelP b -> ModelP (a, b)
pSeq pl pr = pSExpr $ do
  pExact "seq"
  (,) <$> pl <*> pr

  Monadic parser


type VEnv = Map String SExpr

data ModelPState =
  ModelPState {  vEnv   :: VEnv
               , sexps  :: [SExpr]
               }

newtype ModelP a = ModelP { runModelP :: ModelPState -> [(a, ModelPState)] }
  deriving (Functor)

instance Applicative ModelP where
  pure v   = ModelP $ \s -> [(v, s)]
  (<*>)    = ap

instance Monad ModelP where
  list 
    (a, s') <- runModelP m s
    runModelP (f a) s'

instance Alternative ModelP where
  empty    = ModelP $ const []
  m <|> m' = ModelP $ \s -> runModelP m s ++ runModelP m' s

    liftIO m = ModelP $ \s - > m > > = \a - > pure [ ( a , s ) ]

evalModelP' :: ModelP a -> ModelPState -> [a]
evalModelP' m s = map fst (runModelP m s)

evalModelP :: ModelP a -> SExpr -> [a]
evalModelP m sexpr = evalModelP' m (ModelPState Map.empty [sexpr])


pRawHead :: ModelP SExpr
pRawHead = ModelP $ \st -> do
  case sexps st of
    []     -> []
    s : ss -> [(s, st { sexps = ss })]

  Runs the parser p on the first element , taking care of let - related
         ModelP a
pHead pA pS = do
  s <- pRawHead
    r <- handleLet st s
    pure (r, st)
  where
    handleLet st (List [ident -> Just "let", List idents, b]) =
      handleLet (st { vEnv = foldl' addBind (vEnv st) idents }) b
    handleLet st (List [ident -> Just "as", s, _]) = handleLet st s
    handleLet st (Atom a)
      | Just s <- Map.lookup a (vEnv st) = handleLet st s
      | otherwise = evalModelP' (pA a) (st { sexps = [] })
    handleLet st (List ss') = evalModelP' pS (st { sexps = ss' })

    addBind env (List [ident -> Just v, e]) = Map.insert v e env
    addBind _    v = panic "handleLet unexpected let binding" [show v]

pAtom :: ModelP String
pAtom = pHead pure empty

pExact :: String -> ModelP ()
pExact a = pAtom >>= guard . (==) a

pSExpr :: ModelP a -> ModelP a
pSExpr = pHead (const empty)

withSExprs :: [SExpr] -> ModelP a -> ModelP a
withSExprs ss m = ModelP $ \st -> do
  r <- evalModelP' m (st { sexps = ss })
  pure (r, st)

pArray :: Integer -> ModelP a -> ModelP a
pArray count p = do
  (els, arr) <- go
  let v = V.replicate (fromInteger count) arr V.// els
  withSExprs (V.toList v) p
  where
    base :: ModelP ([(Int, SExpr)], SExpr)
    base = pSExpr $ do
      pExact "const"
      s <- pRawHead
      pure ([], s)

    go :: ModelP ([(Int, SExpr)], SExpr)
    go =
      base <|>
      (pSExpr $
        do pExact "store"
           (els, arr) <- go
           i   <- pNumber
           s   <- pRawHead
           pure (if i < count then (fromInteger i, s) : els else els, arr))


   : : String - > ModelP SExpr
   a = do
     < - ModelP $ asks ( Map.lookup a . vEnv )
    case  of
          _ - > panic ( unlines
                  [ " Unexpected response from the SMT solver : "
                  , "   Exptected : a value "



  pTuple : : ModelP a - > ModelP b - > ModelP ( a , b )
  pTuple l r = pSExpr $ do

  pUnit : : ModelP ( )
  pUnit = pExact " unit "


pNumber :: ModelP Integer
pNumber = pAtom >>= go
  where
    go ('#' : 'x' : rest) = do
      let str = "0x" ++ rest
      case readMaybe str of
        Nothing -> error ("Couldn't read " ++ show str)
        Just v  -> pure v
    go ('#' : 'b' : rest) = do
      case binLit rest of
        Nothing -> error ("Couldn't read " ++ show rest)
        Just v  -> pure v
    go ds | all isDigit ds = do
      case readMaybe ds of
        Nothing -> error ("Couldn't read " ++ show ds)
        Just v  -> pure v
    go _ = empty

    binLit cs = do ds <- mapM binDigit cs
                   return $ sum $ zipWith (*) (reverse ds) powers2
    powers2   = 1 : map (2 *) powers2
    binDigit '0' = Just 0
    binDigit '1' = Just 1
    binDigit _   = Nothing

pByte :: ModelP Word8
pByte = fromIntegral <$> pNumber

  This is nicer , we could use  directly
pList :: ModelP a -> ModelP [a]
pList p = ([] <$ pExact "nil") <|>
          (pSExpr $ do pExact "insert"
                       (:) <$> p <*> pList p)

pByteString :: ModelP ByteString
pByteString = BS.pack <$> pList pByte

pBool :: ModelP Bool
pBool = pAtom >>= go
  where
    go "true"  = pure True
    go "false" = pure False
    go _       = empty


pValue :: Type -> ModelP I.Value
pValue ty0 = go ty0
  where
  c.f . symExecTy
    go ty =
      case ty of
        TStream -> unimplemented
        TUInt (TSize n) -> I.vUInt (fromIntegral n) <$> pNumber
        TUInt _ -> unimplemented
        TSInt (TSize n) -> I.vSInt' (fromIntegral n) <$> pNumber
        TSInt _ -> unimplemented
  FIXME
        TBool -> I.VBool <$> pBool
        TUnit -> pure I.vUnit
        TArray _ -> unimplemented
        TMaybe _ty' -> unimplemented
        TMap {} -> unimplemented
        TBuilder {} -> unimplemented
        TIterator {} -> unimplemented
        TUser _ut -> unimplemented
        TParam {} -> unimplemented

    unimplemented = panic "sexprToValue: Unimplemented" [ showPP ty0 ]


ident :: SExpr -> Maybe String
ident (Atom a) = Just a
ident (List [ident -> Just "as", v, _]) = ident v
ident _ = Nothing




   : : Solver - > VEnv - > ModelP a - > String - > IO a
  getValue s e  ( ( Map . ! ? ) e - > Just r ) = runModelP    
  getValue s e  v = do
      _ - > fail $ unlines
                  [ " Unexpected response from the SMT solver : "
                  , "   Exptected : a value "

  pBool : : ModelP 
  pBool = pLisp ( ModelP go )
      go s env ( ident - > Just v ) =  s env  v
      go s env e = error "   

        let str = " 0x " + + rest
        case   of
        case   of
      go s env ( ident - > Just v ) =  s env pNumber v
      go s env e = error " Unparseable byte "

      powers2    = 1 : map ( 2 * ) powers2
      binDigit ' 1 ' = Just 1

   : : ModelP Word8

  pList : : ModelP a - > ModelP [ a ]
          ( ident - > Just v )      - >  s env ( pList p ) v


  pListWithLength p = pLisp ( ModelP go )
          ( ident - > Just v )      - >  s env ( pListWithLength p ) v



  getBytes : : Solver - > SExpr - > IO ByteString
  getBytes s e = BS.pack < $ > runModelP ( pListWithLength  ) s Map.empty e






















  symExecValue : : Type - > I.Value - >   SExpr
  symExecValue   =
    case  of
      VUInt n v   - > S.value ( S.Bits n v )
      VSInt n v   - > S.value ( S.Bits n v )
      VBool b     - > S.value ( S.Bool b )
        mkList ( map ( symExecValue t ) ( Vector.toList vs ) )
       { }     - > unimplemented
      unimplemented = error $ " symExecValue : Unimplemented " + + show 
      mkList : : [ SExpr ] - > SExpr
       = symExecTy 
                  [ foldr ( \a b - > fun " insert " [ a , b ] ) (  ( S.const " nil " ) ( tList  ) ) xs
                  , S.int ( fromIntegral $ length xs )

          TVar x | Just t < - env Map . ! ? x - > go env t
          TVar x - > error ( " Unknown type variable " + + show ( pp x ) )

               { } - > error " Should n't happen ( symExecTy :  ) "
              TStream      - > error " Unimplemented "
              TNum { }      - > error " Should n't happen ( symExecTy : TNum ) "
              TSInt ( Type ( TNum n ) ) - > VSInt ( fromIntegral n ) < $ > pNumber
                - > VBool < $ > pBool
              TArray t '    - > VArray . Vector.fromList < $ > pListWithLength ( sexprToValue t ' )



