{-# LANGUAGE GADTs                #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE PatternGuards        #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE UndecidableInstances #-}

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module    : Yices.Painless.Language
-- Copyright : (c) Don Stewart 2010
-- License   : BSD3
-- Maintainer: Don Stewart <dons@galois.com>
-- Stability : experimental
--
-- This module defines an embedded language of propositions for use by
-- the Yices SMT solver. Propositions are represented as 
-- first-order, typed, pure functions with quantified variables
-- in the proposition encoded as (Haskell) variables bound at the outer scope.
--
-- Propositions are represented as higher-order abstract syntax
-- embedded in Haskell.
-- 
-- Terms and types are embedded in Haskell, so all the usual type
-- inference and checking works for Yices propositions.
--
-- /A simple example: bitvectors/
--
-- > solve $ \b1 (b2 :: Exp BitVector) 
-- >      -> b1 + 1 ==* b2 &&* b2 /=* b2 `xor` 7 + (1 + b1) 
--
-- Constructs the AST term:
--
-- > \x1 x0 -> (&&*) ((==*) ((+) (x1, 0b1), x0),
-- >                  (/=*) (x0, (+) (xor (x0, 0b111), (+) (0b1, x1))))
--
-- And the satisfying assignment:
--
-- > x0 => 0b101
-- > x1 => 0b100
-- > Satisfiable
--

module Yices.Painless.Language (

    -- * Running the solver
    solve,

    -- * Building Yices propositions.
    Yices, Exp,

    -- ** Scalar introduction
    -- constant,
    true, false,

    -- ** Arithmetic.
    -- $Instances

    -- ** Bit vectors
    -- $BitInstances

    -- ** Conditional expressions
    (?),

    -- ** Comparisons.
    (==*), (/=*), (<*), (<=*), (>*), (>=*),

    -- ** Logical operations.
    (&&*), (||*), not,

    -- ** Implication, /n/-ary operations.
    (-->), and, or, max, min,

    -- * Design discussion
    -- $Notes

  ) where

import Prelude hiding (not, or, and, min, max)

import Data.Typeable
import qualified Data.Map as M

import Data.Bits
import qualified Data.Vector.Storable as V

import Control.Monad
import Control.Applicative ((<$>))
import Foreign.Storable (sizeOf)

import Text.PrettyPrint

import qualified Yices.Painless.Base as Yices

import Yices.Painless.Type

import Yices.Painless.Tuple hiding    (Tuple)
import qualified Yices.Painless.Tuple as Tuple

-- $Notes
--
-- /Future work/:
--
-- * Separate 'BitVector's into their own type, so that we can reuse
-- their 'Num' instance, avoiding OverlappingInstances.
--
-- * BitVectors are sized, yet we do not represent their size at all
-- currently.  Size types would let us statically check e.g. bit vector
-- concatenation.
--
-- * BitVectors are represented under the hood as 'Vector Bool'. Yices
-- supports construction also via a 'Word' literal. Support both forms
-- (similar to the 'Integer' type)
--
-- * A 'Monoid' instance for bit vectors?
--
-- * Support other numeric types.
--
-- * Support function types in the core language.
--
-- * Support more Yices types: tuples, records, recursive types.
--


------------------------------------------------------------------------
-- Language

-- We use the dictionary view of overloaded operations (such as arithmetic and
-- bit manipulation) to reify such expressions.  With non-overloaded
-- operations (such as, the logical connectives) and partially overloaded
-- operations (such as comparisons), we use the standard operator names with a
-- '*' attached.  We keep the standard alphanumeric names as they can be
-- easily qualified.

-- Methods from H98 classes, where we need other signatures

infix 4 ==*, /=*,  <*, <=*, >*, >=*

-- | Equality lifted into Yices expressions.
--
(==*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(==*) = mkEq

-- |Inequality lifted into Yices expressions.
--
(/=*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(/=*) = mkNEq

-- | '<' lifted into Yices expressions.
--
(<*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(<*)  = mkLt

-- | '<=' lifted into Yices expressions.
--
(>=*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(>=*) = mkGtEq

-- | '>' lifted into Yices expressions.
--
(>*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(>*)  = mkGt

-- | '<=' lifted into Yices expressions.
--
(<=*) :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
(<=*) = mkLtEq

-- Conditional expressions
-- -----------------------

-- |Conditional expression.
--
infix 0 ?
(?) :: Exp Bool -> (Exp t, Exp t) -> Exp t
c ? (t, e) = Cond c t e

-- Non-overloaded standard functions, where we need other signatures
-- -----------------------------------------------------------------

-- |Conjunction
--
infixr 3 &&*
(&&*) :: Exp Bool -> Exp Bool -> Exp Bool
(&&*) = mkLAnd

-- |Disjunction
--
infixr 2 ||*
(||*) :: Exp Bool -> Exp Bool -> Exp Bool
(||*) = mkLOr

-- | 'and' returns the conjunction of a Boolean list.
and   :: [Exp Bool] -> Exp Bool
and   = foldr (&&*) true

-- | 'or' returns the disjunction of a Boolean list.
or    :: [Exp Bool] -> Exp Bool
or    = foldr (||*) false

-- TODO: implement in terms of underlying primitive conjunction and
-- disjunction, which accepts lists.

-- |Negation
--
not :: Exp Bool -> Exp Bool
not = mkLNot

-- | Implication
infixr 1 -->
(-->) :: Exp Bool -> Exp Bool -> Exp Bool
x --> y = not x ||* y

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

{-
instance (IsScalar t) => Enum (Exp t)
--  succ = mkSucc
--  pred = mkPred
  -- FIXME: ops
-}

instance (IsScalar t) => Prelude.Eq (Exp t) where
  -- FIXME: instance makes no sense with standard signatures
  (==)        = error "Prelude.Eq.== applied to EDSL types"

instance (IsScalar t) => Prelude.Ord (Exp t) where
  -- FIXME: instance makes no sense with standard signatures
  compare     = error "Prelude.Ord.compare applied to EDSL types"

------------------------------------------------------------------------
-- Bit vector operations

-- /TODO/ Vector Bool isn't an instance of num, but abstract Exp BitVectors are!
--
-- We could have an instance of Num BitVector and overload the operators.
-- via the Num (Exp a) instance.
--

instance Num (Exp BitVector) where
    (+)         = mkBVAdd
    (-)         = mkBVSub
    (*)         = mkBVMul
    negate      = mkBVNeg

    fromInteger n = constant . BitVector $
                        V.generate (8 * sizeOf (undefined :: Word)) (testBit w)
        where
            w :: Word
            w = fromIntegral n

    abs _       = error "Prelude.Num.abs applied to EDSL types"
    -- if n >= 0 then n else negate n
    signum _    = error "Prelude.Num.signum applied to EDSL types"

    -- TODO: fromInteger should probably build from mkBVConstant
    -- TOD0: need size information.
    -- TOD0: support construction of a specific size

-- TODO monoid instance
-- instance Monoid (Exp BitVector) where
   --  mappend = mkBVAppend
   -- needs size types!

-- $BitInstances
-- The 'Exp BitVector' type is an instance of 'Bits', allowing the usual
-- Haskell bitwise operators to be used to construct propositions involving
-- bitwise operations on bit vectors.
--
-- > (.&,), (.|.), xor, complement, shiftL, shiftR
--
-- Currently bit vectors are fixed at 'sizeOf (undefined :: Word)' bits.
--
-- Bit vector constants can be constructed using overloaded numeric
-- literals.
--

instance Bits (Exp BitVector) where
    (.&.)       = mkBVAnd
    (.|.)       = mkBVOr
    xor         = mkBVXor
    complement  = mkBVNot
    shiftL      = mkBVShiftL0
    shiftR      = mkBVShiftR0
    isSigned _  = False
    bitSize  _  = 8 * sizeOf (undefined :: Word) -- TODO! size type
        

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

-- $Instances
-- 'Exp' is an instance of 'Num' at the usual types.
--
-- > (+) :: IsNum t => Exp t -> Exp t -> Exp t
--
-- > (-) :: IsNum t => Exp t -> Exp t -> Exp t
--
-- > (*) :: IsNum t => Exp t -> Exp t -> Exp t
--
-- > negate :: IsNum t => Exp t -> Exp t
--
-- Numeric literals are overloaded at Exp t type, so you may write, e.g.
--
-- In addition, 'Exp BitVector' allows for bit vector arithmetic
-- instances.
--

instance (IsNum t) => Num (Exp t) where
  (+)         = mkAdd
  (-)         = mkSub
  (*)         = mkMul

  negate x    = 0 - x

  abs _       = error "Prelude.Num.abs applied to EDSL types"
    -- if n >= 0 then n else negate n
  signum _    = error "Prelude.Num.signum applied to EDSL types"

{-
  signum n | n <  0     = negate 1
           | n == 0     = 0
           | otherwise  = 1
-}

  fromInteger = constant . fromInteger

{-

instance (Elem t, IsNum t, IsIntegral t) => Bits (Exp t) where
    (.&.)      = mkBAnd
    (.|.)      = mkBOr
    xor        = mkBXor
    complement = mkBNot

-}

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

-- |Determine the maximum of two scalars.
max :: IsScalar t => Exp t -> Exp t -> Exp t
max x y = x <* y ? (y , x)

-- |Determine the minimum of two scalars.
min :: IsScalar t => Exp t -> Exp t -> Exp t
min x y = x <* y ? (x , y)

------------------------------------------------------------------------
-- Smart

-- This modules defines the AST of the user-visible embedded language using
-- more convenient higher-order abstract syntax (instead of de Bruijn
-- indices). Moreover, it defines smart constructors to construct programs.

-- Bit vector operators

mkBVAdd :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVAdd x y = PrimBVAdd `PrimApp` tup2 (x, y)

mkBVMul :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVMul x y = PrimBVMul `PrimApp` tup2 (x, y)

mkBVSub :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVSub x y = PrimBVSub `PrimApp` tup2 (x, y)

mkBVNeg :: Exp BitVector -> Exp BitVector
mkBVNeg x = PrimBVNeg `PrimApp` x

mkBVAnd :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVAnd x y = PrimBVAnd `PrimApp` tup2 (x, y)

mkBVOr  :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVOr x y = PrimBVOr `PrimApp` tup2 (x, y)

mkBVXor  :: Exp BitVector -> Exp BitVector -> Exp BitVector
mkBVXor x y = PrimBVXor `PrimApp` tup2 (x, y)

mkBVNot :: Exp BitVector -> Exp BitVector
mkBVNot x = PrimBVNot `PrimApp` x

mkBVShiftL0 :: Exp BitVector -> Int -> Exp BitVector
mkBVShiftL0 x n = PrimBVSL0 `PrimApp` tup2 (x, constant n) -- lift the int shift.

mkBVShiftR0 :: Exp BitVector -> Int -> Exp BitVector
mkBVShiftR0 x n = PrimBVSR0 `PrimApp` tup2 (x, constant n) -- shiftR fills with 0

-- Operators from Num

mkAdd :: (IsNum t) => Exp t -> Exp t -> Exp t
mkAdd x y = PrimAdd numType `PrimApp` tup2 (x, y)

mkSub :: (IsNum t) => Exp t -> Exp t -> Exp t
mkSub x y = PrimSub numType `PrimApp` tup2 (x, y)

mkMul :: (IsNum t) => Exp t -> Exp t -> Exp t
mkMul x y = PrimMul numType `PrimApp` tup2 (x, y)

-- Relational and equality operators

mkLt   :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkLt   x y = PrimLt scalarType `PrimApp` tup2 (x, y)

mkGt   :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkGt   x y = PrimGt scalarType `PrimApp` tup2 (x, y)

mkLtEq :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkLtEq x y = PrimLtEq scalarType `PrimApp` tup2 (x, y)

mkGtEq :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkGtEq x y = PrimGtEq scalarType `PrimApp` tup2 (x, y)

mkEq   :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkEq   x y = PrimEq  scalarType `PrimApp` tup2 (x, y)

mkNEq  :: (IsScalar t) => Exp t -> Exp t -> Exp Bool
mkNEq  x y = PrimNEq scalarType `PrimApp` tup2 (x, y)

-- Logical operators

mkLAnd :: Exp Bool -> Exp Bool -> Exp Bool
mkLAnd x y = PrimLAnd `PrimApp` tup2 (x, y)

mkLOr :: Exp Bool -> Exp Bool -> Exp Bool
mkLOr x y = PrimLOr `PrimApp` tup2 (x, y)

mkLNot :: Exp Bool -> Exp Bool
mkLNot x = PrimLNot `PrimApp` x

-- Smart constructor for literals

-- | Literal 'True'
true :: Exp Bool
true  = constant True

-- | Literal 'False'
false :: Exp Bool
false = constant False

-- |Constant scalar expression
--
constant :: (Show t, IsScalar t) => t -> Exp t
constant = Const

-- Smart constructor and destructors for tuples

tup2 :: (Exp a, Exp b) -> Exp (a, b)
tup2 (x1, x2) = Tuple (NilTup `SnocTup` x1 `SnocTup` x2)

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

-- Embedded expressions of the surface language

-- HOAS expressions mirror the constructors of `OpenExp', but with the
-- `Tag' constructor instead of variables in the form of de Bruijn indices.
-- Moreover, HOAS expression use n-tuples and the type class 'Elem' to
-- constrain element types, whereas `OpenExp' uses nested pairs and the 
-- GADT 'TupleType'.

--
-- Yices programs should be closed wrt. scalar variables.
-- All /variables/ must be 'defined' in Yices before use.
-- Nesting should be ok. So bound under a lambda.
--
-- To restrict functions to first-order, we separate function
-- abstraction from the main expression type.  Functions are represented
-- using de Bruijn indices.

-- Helpful wrapper
data YicesProg r where
    Y :: Yices f r => f -> YicesProg r

--
-- Allow embedded functions??

-- | Yices scalar formula.
--
data Exp t where

    -- Needed for conversion to de Bruijn form
    Tag         :: (IsScalar t) => Int          ->  Exp t
                 -- environment size at defining occurrence

    Const       :: (Show t, IsScalar t)
                => t                             -> Exp t

    Tuple       :: (IsTuple t)
                => Tuple.Tuple Exp (TupleRepr t) -> Exp t

    Cond        :: Exp Bool -> Exp t -> Exp t    -> Exp t

    PrimApp     :: PrimFun (a -> r) -> Exp a     -> Exp r

{-

  Lam :: (Typeable s, Typeable t,
          Show s, Show t)
      => (Term s -> Term t)      -> Term (s -> t)

  App :: (Typeable s, Typeable t,
          Show s, Show t)
      => Term (s -> t) -> Term s -> Term t

 -}

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

-- | Well-formed Yices propositions (with /n/ quantified variables).
--
-- The current design requires all free variables in the proposition to
-- be bound at the outermost level. In higher-order abstract syntax,
-- this represents an n-ary, polyvariadic function.
--
-- Examples:
--
-- > true
--
-- > \(x :: Exp Int) -> x >* 8 &&* x <* 10
--
-- > \x y -> x ==* y + 1
--
-- The language supports polymorphic scalar and numerical operations, so
-- some explicit type information may be required to resolve overloading.
-- 
-- E.g.
-- 
-- > > solve $ \x y -> x ==* y + (1 :: Exp Int) 
-- >
-- > \x1 x0 -> (==*) (x1, (+) (x0, 1))
-- >
-- > x0 => 1
-- > x1 => 2
-- > Satisfiable
--
class Yices f r | f -> r where
  -- Convert a HOAS fragment into its deBruijn form, binding variables in a typed environment
  convert :: Layout env env -> f -> OpenFun env r

instance Yices (Exp b) b where
  -- Convert expressions
  convert lyt e = OBody (convertOpenExp lyt e)

instance (IsScalar a, Yices f r) => Yices (Exp a -> f) (a -> r) where
  -- Convert binders, one bind at a time.
  convert lyt f = OLam (convert lyt' (f a))
    where
    a    = Tag (size lyt)
    lyt' = inc lyt `PushLayout` ZeroIdx

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

-- |Conversion from HOAS to de Bruijn expression AST
--
-- Based on:
--
-- * Chakravarty. /Converting a HOAS term GADT into a de Bruijn term GADT/, 2009.
-- <http://www.cse.unsw.edu.au/~chak/haskell/term-conv/>
--

-- |Convert a closed Yices program.
--
convertYices :: YicesProg r -> OpenYices r
convertYices (Y f) = OY $ convert EmptyLayout f

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

-- |Convert an open expression with a given environment layout.
--             
convertOpenExp :: forall t env.
                  Layout env  env       -- scalar environment 
               -> Exp t                 -- expression to be converted
               -> OpenExp env t
convertOpenExp lyt = cvt
  where
    cvt :: Exp t' -> OpenExp env t'
    cvt (Tag i)             = Var (prjIdx (size lyt - i - 1) lyt) -- indexing!
    cvt (Const v)           = OConst v
    cvt (Tuple tup)         = OTuple (convertTuple lyt tup)
    cvt (Cond e1 e2 e3)     = OCond (cvt e1) (cvt e2) (cvt e3)
    cvt (PrimApp p e)       = OPrimApp p (cvt e)

-- |Convert a tuple expression
--
convertTuple :: Layout env env
             -> Tuple.Tuple Exp t
             -> Tuple.Tuple (OpenExp env) t
convertTuple _lyt NilTup           = NilTup
convertTuple lyt  (es `SnocTup` e) = convertTuple lyt es `SnocTup` convertOpenExp lyt e

-- |Convert an expression closed wrt to scalar variables
--
convertExp :: Exp t -> OExp t
convertExp = convertOpenExp EmptyLayout

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

-- |
-- A layout of an environment an entry for each entry of the environment.
-- Each entry in the layout holds the deBruijn index that refers to the
-- corresponding entry in the environment.
--
-- TODO: explain the two type variables
--
data Layout env env' where
  EmptyLayout :: Layout env ()
  PushLayout  :: Typeable t
              => Layout env env' -> Idx env t -> Layout env (env', t)

-- Project the nth index out of an environment layout.
--
prjIdx :: forall t env env'. Typeable t => Int -> Layout env env' -> Idx env t
prjIdx 0 (PushLayout _ (ix :: Idx env u)) = case gcast ix of
   Just ix' -> ix'
   Nothing  -> error $
                "EDSL Compiler Type Error.\n" ++
                "Couldn't match expected type `" ++ show (typeOf (undefined :: t)) ++
                     "' against inferred type `" ++ show (typeOf (undefined :: u)) ++ "'"

prjIdx n (PushLayout l _)  = prjIdx (n - 1) l
prjIdx _ EmptyLayout       =
  error "Yices.Painless.Language.prjIdx" "inconsistent valuation"

-- | More functions on layouts, from 
-- <http://www.cse.unsw.edu.au/~chak/haskell/term-conv/Convert.hs>
--

-- Yield the number of entries in an environment layout
--
size :: Layout env env' -> Int
size EmptyLayout        = 0
size (PushLayout lyt _) = size lyt + 1

-- Add an entry to a layout, incrementing all indices
--
inc :: Layout env env' -> Layout (env, t) env'
inc EmptyLayout         = EmptyLayout
inc (PushLayout lyt ix) = PushLayout (inc lyt) (SuccIdx ix)

------------------------------------------------------------------------
-- AST 

--
-- The embedded language is marginally two-levels. Formula with free variables 
-- are bound with Define. Thus Yices programs are closed expressions of yices formula.
--
-- There is no explicit sharing in the initial AST form, but sharing can be recovered
-- subsequently by common subexpression elimination.
--
-- TODO: add functions to the Exp type.
--
-- The AST contains both reified dictionaries and type class constraints.  
-- Type classes are used for yices-related functionality that is uniformly
-- available for all supported types.  In contrast, reified dictionaries are
-- used for functionality that is only available for certain types, such as
-- arithmetic operations.
--

-- | Yices computations parameterized by Yices /variables/ represented with de
-- Bruijn indices.
--
-- * Scalar functions and expressions embedded in well-formed Yices
--   computations cannot contain free scalar variable indices.  The latter
--   cannot be bound in Yices computations, and hence, cannot appear in any
--   well-formed program.
--
-- The data type is parametrised over the surface types (not the representation
-- type).

-- data OpenYices a where
--    ODefine :: IsScalar t => Fun (t -> u) -> OpenYices u

-- | Yices programs with an environment.
--        
data OpenYices t where
    OY :: OFun t -> OpenYices t

-- execF shows nicely how to recurse over the OpenFun structure.
  
-- |Function abstraction
--
data OpenFun env t where
  OBody ::               OpenExp env      t -> OpenFun env t
  OLam  :: IsScalar a => OpenFun (env, a) t -> OpenFun env (a -> t)

-- |Function without free scalar variables
--
type OFun t = OpenFun () t

-- |Open expressions using de Bruijn indices for variables ranging over tuples
-- of scalars.  All code, except Cond, is evaluated eagerly.  N-tuples are
-- represented as nested pairs. 
--
-- The data type is parametrised over the surface types (not the representation
-- type).
--
data OpenExp env t where

  -- Variable index, ranging only over tuples or scalars
  Var         :: IsScalar t
              => Idx env t
              -> OpenExp env t

  -- Constant values
  OConst      :: (Show t, IsScalar t)
              => t
              -> OpenExp env t

  -- Tuples   
  OTuple      :: (IsTuple t)
              => Tuple.Tuple (OpenExp env) (TupleRepr t)
              -> OpenExp env t

  -- Conditional expression
  OCond       :: OpenExp env Bool
              -> OpenExp env t
              -> OpenExp env t
              -> OpenExp env t

  -- Primitive scalar operations
  OPrimApp     :: PrimFun (a -> r)
               -> OpenExp env a
               -> OpenExp env r


-- |Expression without free scalar variables
--
type OExp t = OpenExp () t

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

-- Typed de Bruijn indices
-- -----------------------

-- De Bruijn variable index projecting a specific type from a type
-- environment.  Type envionments are nested pairs (..((), t1), t2, ..., tn). 
--
data Idx env t where
  ZeroIdx ::              Idx (env, t) t
  SuccIdx :: Idx env t -> Idx (env, s) t

-- Environments
-- ------------

-- Valuation for an environment
--
{-
data Val env where
--  Empty :: Val ()
--  Push  :: Val env -> t -> Val (env, t)

-- Projection of a value from a valuation using a de Bruijn index
--
prj :: Idx env t -> Val env -> t
prj ZeroIdx       (Push _   v) = v
prj (SuccIdx idx) (Push val _) = prj idx val
prj _             _            =
  error "Yices.Painless.Language.prj" "inconsistent valuation"
-}

-- Convert a typed de Brujin index to the corresponding integer
--
idxToInt :: Idx env t -> Int
idxToInt ZeroIdx       = 0
idxToInt (SuccIdx idx) = 1 + idxToInt idx

------------------------------------------------------------------------
-- AST

-- |Primitive scalar operations
--
data PrimFun sig where
    -- relational and equality operators
    PrimLt   :: ScalarType a -> PrimFun ((a, a) -> Bool)
    PrimGt   :: ScalarType a -> PrimFun ((a, a) -> Bool)
    PrimLtEq :: ScalarType a -> PrimFun ((a, a) -> Bool)
    PrimGtEq :: ScalarType a -> PrimFun ((a, a) -> Bool)
    PrimEq   :: ScalarType a -> PrimFun ((a, a) -> Bool)
    PrimNEq  :: ScalarType a -> PrimFun ((a, a) -> Bool)

    -- operators from Num
    PrimAdd  :: NumType a -> PrimFun ((a, a) -> a)
    PrimSub  :: NumType a -> PrimFun ((a, a) -> a)
    PrimMul  :: NumType a -> PrimFun ((a, a) -> a)

    -- logical operators
    PrimLAnd :: PrimFun ((Bool, Bool) -> Bool)
    PrimLOr  :: PrimFun ((Bool, Bool) -> Bool)
    PrimLNot :: PrimFun (Bool         -> Bool)

    -- TODO: just use the overloaded instance?

    -- bit vector arithmetic
    PrimBVAdd :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVSub :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVMul :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVNeg :: PrimFun (BitVector -> BitVector)

    -- bit vector bit operations
    PrimBVAnd :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVOr  :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVXor :: PrimFun ((BitVector,BitVector) -> BitVector)
    PrimBVNot :: PrimFun (BitVector -> BitVector)
    -- Todo: enforce size constraints.

    PrimBVSL0 :: PrimFun ((BitVector,Int) -> BitVector)
    PrimBVSR0 :: PrimFun ((BitVector,Int) -> BitVector)

    -- Others are already overloaded on Eq/Neq/Lt/Lte/Ge/Gte

-- Bit vector logical operators
{-
    mkBVConcat, mkBVExtract,
    mkBVSignExtend,
    mkBVShiftLeft1,
    mkBVShiftRight1,
    mkBVSlt, mkBVSle, -- can be interpreted as signed types.
    mkBVSgt, mkBVSge,
-}

------------------------------------------------------------------------
-- Pretty printing

-- Generic instance for the AST, to avoid a wrapper type.
instance Yices f r => Show f where
    show = show . Y

instance Show (YicesProg as) where
  show = show . convertYices

instance Show (Exp a) where
  show = show . convertExp

instance Show (OpenYices a) where
  show c = render $ prettyYices 0 c

instance Show (OpenFun env f) where
  show f = render $ prettyFun 0 f

instance Show (OpenExp env t) where
  show e = render $ prettyExp 0 noParens e

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

prettyYices  :: Int -> OpenYices a -> Doc
prettyYices n (OY f) = prettyFun n f

prettyFun ::  Int -> OpenFun env a -> Doc
prettyFun lvl fun =
  let (n, bodyDoc) = count fun
  in
    if n < 0 
        then bodyDoc
        else 
          char '\\' <> hsep [text $ "x" ++ show idx | idx <- reverse [0..n]] <+>
          text "->" <+> bodyDoc
  where
     count :: OpenFun env fun -> (Int, Doc)
     count (OBody body) = (-1, prettyExp lvl noParens body)
     count (OLam fun')  = let (n, body) = count fun' in (1 + n, body)

-- Pretty print an expression.
--
-- * Apply the wrapping combinator (1st argument) to any compound expressions.
--
prettyExp :: forall t env .
             Int -> (Doc -> Doc) -> OpenExp env t -> Doc

prettyExp _   _    (Var idx)         = text $ "x" ++ show (idxToInt idx)
prettyExp _   _    (OConst v)        = text $ show (v :: t) -- dispatch differently for BitVector types
prettyExp lvl _    (OTuple tup)      = prettyTuple lvl tup

-- prettyExp lvl wrap (Prj idx e)
--  = wrap $ prettyTupleIdx idx <+> prettyExp lvl parens e

prettyExp lvl wrap (OCond c t e)
     = wrap $ sep [prettyExp lvl parens c <+> char '?',
                parens (prettyExp lvl noParens t <> comma <+>
                        prettyExp lvl noParens e)]

prettyExp lvl wrap (OPrimApp p a)
    = wrap $ prettyPrim p <+> prettyExp lvl parens a

-- Pretty print nested pairs as a proper tuple.
--
prettyTuple :: Int -> Tuple.Tuple (OpenExp env) t -> Doc
prettyTuple lvl e = parens $ sep (map (<> comma) (init es) ++ [last es])
  where
    es = collect e

    collect :: Tuple.Tuple (OpenExp env) t -> [Doc]
    collect NilTup          = []
    collect (SnocTup tup e') = collect tup ++ [prettyExp lvl noParens e']
    
-- Pretty print a primitive operation
--
prettyPrim :: PrimFun a -> Doc
prettyPrim (PrimAdd _)         = text "(+)"
prettyPrim (PrimSub _)         = text "(-)"
prettyPrim (PrimMul _)         = text "(*)"
prettyPrim (PrimLt _)          = text "(<*)"
prettyPrim (PrimGt _)          = text "(>*)"
prettyPrim (PrimLtEq _)        = text "(<=*)"
prettyPrim (PrimGtEq _)        = text "(>=*)"
prettyPrim (PrimEq _)          = text "(==*)"
prettyPrim (PrimNEq _)         = text "(/=*)"
prettyPrim PrimLAnd            = text "(&&*)"
prettyPrim PrimLOr             = text "(||*)"
prettyPrim PrimLNot            = text "not"

-- Overloaded bit vector ops
prettyPrim PrimBVAdd           = text "(+)"
prettyPrim PrimBVMul           = text "(*)"
prettyPrim PrimBVSub           = text "(-)"
prettyPrim PrimBVNeg           = text "negate"
prettyPrim PrimBVAnd           = text "(.&.)"
prettyPrim PrimBVOr            = text "(.|.)"
prettyPrim PrimBVXor           = text "xor"
prettyPrim PrimBVNot           = text "complement"
prettyPrim PrimBVSL0           = text "shiftL"
prettyPrim PrimBVSR0           = text "shiftR"

-- Auxilliary pretty printing combinators
-- 
noParens :: Doc -> Doc
noParens = id


------------------------------------------------------------------------
--
-- Execute the AST.
--
-- TODO: check use of 'Context' is pure.
--

-- | Run Yices on a well-formed proposition, returning either a
-- satisfying assignment of variables that makes the proposition hold, or
-- establish that the proposition is unsatisfiable.
--
solve :: (Yices f r) => f -> IO Yices.Result
solve q' = do
    let q = Y q'
    c <- Yices.mkContext 
    Yices.setTypeChecker True

    let t = convertYices q -- bind all the variables
    print t

    (g,e) <- execY c t 

    Yices.assert c e

--    Yices.ctxDump c
--    Yices.ppExpr e

    sat <- Yices.check c
    case sat of
        Yices.Unsatisfiable -> return $ Yices.Unsatisfiable
        Yices.Undefined     -> return $ Yices.Undefined
        Yices.Satisfiable   -> do
            mm <- Yices.getModel c 
            case mm of
                Nothing -> return Yices.Satisfiable
                Just m  -> do

                    vs <- sequence
                            [ get m v t' d
                            | (v,(t',d)) <- M.toList g ]

                    forM_ vs (\(n, YValue mv) -> do
                                putStr (n ++ " => ")
                                case mv of
                                    Nothing -> putStrLn "_"
                                    Just v  -> print v)

                    -- print cs
                    return Yices.Satisfiable

-- | Retrieving bindings by type
--
get :: Yices.Model -> String -> YType -> Yices.Decl -> IO (String, YValue)
get m v (YType (ty :: ScalarType t)) d

    | NumScalarType (IntegralNumType (TypeInt _))  <- ty
    = do mn <- Yices.getValueInt m d
         return (v, YValue mn)

    | NonNumScalarType (TypeBool _)                <- ty
    = do mn <- Yices.getValueBool m d
         return (v, YValue mn)

    | NonNumScalarType (TypeVectorBool _)                <- ty
    = do mn <- Yices.getValueBitVector m d (fromIntegral $ 8 * sizeOf (undefined :: Word))
         return (v, YValue $ fmap BitVector mn)
            
    | otherwise = error "Yices.Painless.get: don't know how to get this type yet"


-- | Map free variable names that we come up with to their type and decl
-- in the context. This will then be used to extract solutions after
-- solving.
--
type YEnv = M.Map String (YType, Yices.Decl)

data YType  = forall a. IsScalar a => YType (ScalarType a)

data YValue = forall a. (Show a, IsScalar a) => YValue (Maybe a)

--
-- To run,
--
--   * take the Exp
--   * extract the set of declared variables
--   * bind them in a context
--   * evaluate the assertions.
-- 
-- return the bindings.
--
execY :: Yices.Context -> OpenYices t -> IO (YEnv, Yices.Expr)
execY c (OY f) = execF c f

--
-- Declaring variables. Begin with a closed Yices program.
--
execF :: Yices.Context -> OFun t -> IO (YEnv, Yices.Expr)
execF c fn = go fn 0 M.empty
    where
        go :: OpenFun env u -> Int -> YEnv -> IO (YEnv, Yices.Expr)
        go (OBody b) _ g = (,) g <$> exec c b

        -- Numbers
        go (OLam (f :: OpenFun (env, a) t)) n g
            | ty@(NumScalarType (IntegralNumType (TypeInt _)))  <- scalarType :: ScalarType a
            = do
                let nm = "x" ++ show n -- we get to actually name things
                tynm <- Yices.mkType c "int"
                d    <- Yices.mkVarDecl c nm tynm
                go f (n + 1) (M.insert nm (YType ty, d) g)
        
        -- Booleans
        go (OLam (f :: OpenFun (env, a) t)) n g
            | ty@(NonNumScalarType (TypeBool _))                <- scalarType :: ScalarType a
            = do
                let nm = "x" ++ show n
                d <- Yices.mkBoolDecl c nm
                go f (n + 1) (M.insert nm (YType ty, d) g)

        -- Bit vectors
        go (OLam (f :: OpenFun (env, a) t)) n g
            | ty@(NonNumScalarType (TypeVectorBool _))                <- scalarType :: ScalarType a
            = do
                let nm = "x" ++ show n
                -- TODO: hack, no size information for bit vectors yet.
                -- TODO: show bvs in a better form.
                tynm <- Yices.mkBitVectorType c (fromIntegral $ 8 * sizeOf (undefined :: Word))
                d    <- Yices.mkVarDecl c nm tynm
                go f (n + 1) (M.insert nm (YType ty, d) g)


        go _ _ _ = error "Yices.execF: don't know how to bind variables of this type yet"

-- | Execute an expression with free variables.
--
-- /TODO:/
--
-- * support bit vector operations.
-- 
-- * function application and creation other than at the top level
--
-- * Rational, Double types.
--
-- * Tuples
--
exec :: forall env t. Yices.Context -> OpenExp env t -> IO Yices.Expr

-- GADTs are magic
exec c (OConst n)
    | NumScalarType (IntegralNumType (TypeInt _))   <- scalarType :: ScalarType t
    = Yices.mkNum c n

    | NonNumScalarType (TypeBool _)                 <- scalarType :: ScalarType t
    = if n then Yices.mkTrue  c
           else Yices.mkFalse c

    | NonNumScalarType (TypeVectorBool _)            <- scalarType :: ScalarType t
    = Yices.mkBVConstantFromVector c (unBV n)


exec c (Var i) = do
    let n = "x" ++ show (idxToInt i)
--    print n

    v <- Yices.getVarDeclFromName c n -- sneaky. using Yices environment. TODO: use YEnv
    case v of
        Nothing -> error "Undefined variable"
        Just d  -> Yices.mkVarFromDecl c d

-- Conditionals
exec c (OCond b t e) = do
            eb <- exec c b
            et <- exec c t
            ee <- exec c e
            Yices.mkIte c eb et ee

-- Works on bitvectors straight up.
exec c (OPrimApp  (PrimEq _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkEq c e1 e2

-- Ok for bitvectors
exec c (OPrimApp  (PrimNEq _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkNeq c e1 e2

-- Specialize for BitVector first
exec c (OPrimApp  (PrimLt (NonNumScalarType (TypeVectorBool _)))
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVLt c e1 e2

-- Only numeric types?
exec c (OPrimApp  (PrimLt _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkLt c e1 e2

-- Specialize for BitVector first
exec c (OPrimApp  (PrimLtEq (NonNumScalarType (TypeVectorBool _)))
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVLe c e1 e2

exec c (OPrimApp  (PrimLtEq _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkLe c e1 e2

-- Specialize for BitVector first
exec c (OPrimApp  (PrimGt (NonNumScalarType (TypeVectorBool _)))
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVGt c e1 e2

exec c (OPrimApp  (PrimGt _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkGt c e1 e2

-- TODO think about what gets through the overloading incorrectly.
-- Specialize for BitVector first
exec c (OPrimApp  (PrimGtEq (NonNumScalarType (TypeVectorBool _)))
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVGe c e1 e2

exec c (OPrimApp  (PrimGtEq _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkGe c e1 e2

exec c (OPrimApp  PrimLAnd
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkAnd c [e1, e2]

exec c (OPrimApp  PrimLOr
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkOr c [e1, e2]

exec c (OPrimApp PrimLNot x) = do
            e <- exec c x
            Yices.mkNot c e

-- Numerical operations

exec c (OPrimApp  (PrimAdd _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkSum c [e1,e2]

exec c (OPrimApp  (PrimSub _)
        (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkSub c [e1,e2]

exec c (OPrimApp  (PrimMul _)
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkMul c [e1,e2]

------------------------------------------------------------------------
-- Bit Vector operations

exec c (OPrimApp PrimBVAdd
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVAdd c e1 e2

exec c (OPrimApp PrimBVMul
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVMul c e1 e2

exec c (OPrimApp PrimBVSub
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVSub c e1 e2

exec c (OPrimApp PrimBVNeg
          (OTuple (NilTup `SnocTup` x1))) = do
            e1 <- exec c x1
            Yices.mkBVMinus c e1

exec c (OPrimApp PrimBVAnd
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVAnd c e1 e2

exec c (OPrimApp PrimBVOr
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVOr c e1 e2

exec c (OPrimApp PrimBVXor
          (OTuple (NilTup `SnocTup` x1 `SnocTup` x2))) = do
            e1 <- exec c x1
            e2 <- exec c x2
            Yices.mkBVXor c e1 e2

exec c (OPrimApp PrimBVNot
          (OTuple (NilTup `SnocTup` x1))) = do
            e1 <- exec c x1
            Yices.mkBVNot c e1

exec c (OPrimApp PrimBVSL0
          (OTuple (NilTup `SnocTup` x1 `SnocTup` (OConst n)))) = do
            e1 <- exec c x1
            Yices.mkBVShiftLeft0 c e1 n

exec c (OPrimApp PrimBVSR0
          (OTuple (NilTup `SnocTup` x1 `SnocTup` (OConst n)))) = do
            e1 <- exec c x1
            Yices.mkBVShiftRight0 c e1 n

exec _ _ = error "Not implemented"

{-
 _ (Tuple _)
 _ (Const _)
 _ (PrimApp (PrimEq _) (Const _))
-}

-- exec _ e = error (show e)

{-
evalTuple :: Tuple (OpenExp env aenv) t -> Val env -> Val aenv -> t
evalTuple NilTup            _env _aenv = ()
evalTuple (tup `SnocTup` e) env  aenv  = (evalTuple tup env aenv, 
                                          evalOpenExp e env aenv)
-}