{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE OverloadedStrings    #-}

{-# OPTIONS_GHC -Wno-name-shadowing #-}

--------------------------------------------------------------------------------
-- | `defunctionalize` traverses the query to:
--      1. "normalize" lambda terms by renaming binders,
--      2. generate alpha- and beta-equality axioms for
--   The lambdas and redexes found in the query.
--
--   NOTE: `defunctionalize` should happen **BEFORE**
--   `elaborate` as the latter converts all actual `EApp`
--   into the (uninterpreted) `smt_apply`.
--   We cannot elaborate prior to `defunc` as we need the
--   `EApp` and `ELam` to determine the lambdas and redexes.
--------------------------------------------------------------------------------

module Language.Fixpoint.Defunctionalize
  ( defunctionalize
  , Defunc(..)
  , defuncAny
  ) where

import qualified Data.HashMap.Strict as M
import           Data.Hashable
import           Control.Monad.State
import           Language.Fixpoint.Misc            (fM, secondM, mapSnd)
import           Language.Fixpoint.Solver.Sanitize (symbolEnv)
import           Language.Fixpoint.Types        hiding (GInfo(..), allowHO, fi)
import qualified Language.Fixpoint.Types           as Types (GInfo(..))
import           Language.Fixpoint.Types.Config
import           Language.Fixpoint.Types.Visitor   (mapMExpr)
-- import Debug.Trace (trace)

defunctionalize :: (Fixpoint a) => Config -> SInfo a -> SInfo a
defunctionalize :: forall a. Fixpoint a => Config -> SInfo a -> SInfo a
defunctionalize Config
cfg SInfo a
si = forall s a. State s a -> s -> a
evalState (forall a. Defunc a => a -> DF a
defunc SInfo a
si) (forall a. Config -> SInfo a -> DFST
makeInitDFState Config
cfg SInfo a
si)

defuncAny :: Defunc a => Config -> SymEnv -> a -> a
defuncAny :: forall a. Defunc a => Config -> SymEnv -> a -> a
defuncAny Config
cfg SymEnv
env a
e = forall s a. State s a -> s -> a
evalState (forall a. Defunc a => a -> DF a
defunc a
e) (Config -> SymEnv -> IBindEnv -> DFST
makeDFState Config
cfg SymEnv
env IBindEnv
emptyIBindEnv)


---------------------------------------------------------------------------------------------
-- | Expressions defunctionalization --------------------------------------------------------
---------------------------------------------------------------------------------------------
txExpr :: Expr -> DF Expr
txExpr :: Expr -> DF Expr
txExpr Expr
e = do
  Bool
hoFlag <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets DFST -> Bool
dfHO
  if Bool
hoFlag then Expr -> DF Expr
defuncExpr Expr
e else forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e

defuncExpr :: Expr -> DF Expr
defuncExpr :: Expr -> DF Expr
defuncExpr = forall (m :: * -> *). Monad m => (Expr -> m Expr) -> Expr -> m Expr
mapMExpr Expr -> DF Expr
reBind
         forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> forall (m :: * -> *). Monad m => (Expr -> m Expr) -> Expr -> m Expr
mapMExpr (forall (m :: * -> *) a b. Monad m => (a -> b) -> a -> m b
fM Expr -> Expr
normalizeLams)

reBind :: Expr -> DF Expr
reBind :: Expr -> DF Expr
reBind (ELam (Symbol
x, Sort
s) Expr
e) = (\Symbol
y -> (Symbol, Sort) -> Expr -> Expr
ELam (Symbol
y, Sort
s) (forall a. Subable a => a -> (Symbol, Expr) -> a
subst1 Expr
e (Symbol
x, Symbol -> Expr
EVar Symbol
y))) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sort -> DF Symbol
freshSym Sort
s
reBind Expr
e               = forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e
shiftLam :: Int -> Symbol -> Sort -> Expr -> Expr
shiftLam :: Int -> Symbol -> Sort -> Expr -> Expr
shiftLam Int
i Symbol
x Sort
t Expr
e = (Symbol, Sort) -> Expr -> Expr
ELam (Symbol
x_i, Sort
t) (Expr
e forall a. Subable a => a -> (Symbol, Expr) -> a
`subst1` (Symbol
x, Expr
x_i_t))
  where
    x_i :: Symbol
x_i          = Int -> Symbol
lamArgSymbol Int
i
    x_i_t :: Expr
x_i_t        = Expr -> Sort -> Expr
ECst (Symbol -> Expr
EVar Symbol
x_i) Sort
t

-- | normalize lambda arguments [TODO: example]
--
-- Renames lambda bindings to lamb_arg##i. Each use of a lambda binding
-- is surrounded with a cast.

normalizeLams :: Expr -> Expr
normalizeLams :: Expr -> Expr
normalizeLams Expr
e = forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Int -> Expr -> (Int, Expr)
normalizeLamsFromTo Int
1 Expr
e

normalizeLamsFromTo :: Int -> Expr -> (Int, Expr)
normalizeLamsFromTo :: Int -> Expr -> (Int, Expr)
normalizeLamsFromTo Int
i   = Expr -> (Int, Expr)
go
  where
    go :: Expr -> (Int, Expr)
go (ELam (Symbol
y, Sort
sy) Expr
e) = (Int
i' forall a. Num a => a -> a -> a
+ Int
1, Int -> Symbol -> Sort -> Expr -> Expr
shiftLam Int
i' Symbol
y Sort
sy Expr
e') where (Int
i', Expr
e') = Expr -> (Int, Expr)
go Expr
e
                          -- let (i', e') = go e
                          --    y'       = lamArgSymbol i'  -- SHIFTLAM
                          -- in (i' + 1, ELam (y', sy) (e' `subst1` (y, EVar y')))
    go (EApp Expr
e1 Expr
e2)     = let (Int
i1, Expr
e1') = Expr -> (Int, Expr)
go Expr
e1
                              (Int
i2, Expr
e2') = Expr -> (Int, Expr)
go Expr
e2
                          in (forall a. Ord a => a -> a -> a
max Int
i1 Int
i2, Expr -> Expr -> Expr
EApp Expr
e1' Expr
e2')
    go (ECst Expr
e Sort
s)       = forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd (Expr -> Sort -> Expr
`ECst` Sort
s) (Expr -> (Int, Expr)
go Expr
e)
    go (PAll [(Symbol, Sort)]
bs Expr
e)      = forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd ([(Symbol, Sort)] -> Expr -> Expr
PAll [(Symbol, Sort)]
bs) (Expr -> (Int, Expr)
go Expr
e)
    go Expr
e                = (Int
i, Expr
e)


--------------------------------------------------------------------------------
-- | Containers defunctionalization --------------------------------------------
--------------------------------------------------------------------------------

class Defunc a where
  defunc :: a -> DF a

instance (Defunc (c a), TaggedC c a) => Defunc (Types.GInfo c a) where
  defunc :: GInfo c a -> DF (GInfo c a)
defunc GInfo c a
fi = do
    HashMap SubcId (c a)
cm'    <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
Types.cm    GInfo c a
fi
    HashMap KVar (WfC a)
ws'    <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> HashMap KVar (WfC a)
Types.ws    GInfo c a
fi
    -- NOPROP setBinds $ mconcat ((senv <$> M.elems (cm fi)) ++ (wenv <$> M.elems (ws fi)))
    SEnv Sort
gLits' <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> SEnv Sort
Types.gLits GInfo c a
fi
    SEnv Sort
dLits' <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> SEnv Sort
Types.dLits GInfo c a
fi
    BindEnv a
bs'    <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> BindEnv a
Types.bs    GInfo c a
fi
    [Triggered Expr]
ass'   <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> [Triggered Expr]
Types.asserts GInfo c a
fi
    -- NOPROP quals' <- defunc $ quals fi
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ GInfo c a
fi { cm :: HashMap SubcId (c a)
Types.cm      = HashMap SubcId (c a)
cm'
                , ws :: HashMap KVar (WfC a)
Types.ws      = HashMap KVar (WfC a)
ws'
                , gLits :: SEnv Sort
Types.gLits   = SEnv Sort
gLits'
                , dLits :: SEnv Sort
Types.dLits   = SEnv Sort
dLits'
                , bs :: BindEnv a
Types.bs      = BindEnv a
bs'
                , asserts :: [Triggered Expr]
Types.asserts = [Triggered Expr]
ass'
                }

instance (Defunc a) => Defunc (Triggered a) where
  defunc :: Triggered a -> DF (Triggered a)
defunc (TR Trigger
t a
e) = forall a. Trigger -> a -> Triggered a
TR Trigger
t forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc a
e

instance Defunc (SimpC a) where
  defunc :: SimpC a -> DF (SimpC a)
defunc SimpC a
sc = do Expr
crhs' <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall a. SimpC a -> Expr
_crhs SimpC a
sc
                 forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ SimpC a
sc {_crhs :: Expr
_crhs = Expr
crhs'}

instance Defunc (WfC a) where
  defunc :: WfC a -> DF (WfC a)
defunc wf :: WfC a
wf@WfC{} = do
    let (Symbol
x, Sort
t, KVar
k) = forall a. WfC a -> (Symbol, Sort, KVar)
wrft WfC a
wf
    Sort
t' <- forall a. Defunc a => a -> DF a
defunc Sort
t
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ WfC a
wf { wrft :: (Symbol, Sort, KVar)
wrft = (Symbol
x, Sort
t', KVar
k) }
  defunc wf :: WfC a
wf@GWfC{} = do
    let (Symbol
x, Sort
t, KVar
k) = forall a. WfC a -> (Symbol, Sort, KVar)
wrft WfC a
wf
    Sort
t' <- forall a. Defunc a => a -> DF a
defunc Sort
t
    Expr
e' <- forall a. Defunc a => a -> DF a
defunc forall a b. (a -> b) -> a -> b
$ forall a. WfC a -> Expr
wexpr WfC a
wf
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ WfC a
wf { wrft :: (Symbol, Sort, KVar)
wrft = (Symbol
x, Sort
t', KVar
k), wexpr :: Expr
wexpr = Expr
e' }

instance Defunc SortedReft where
  defunc :: SortedReft -> DF SortedReft
defunc (RR Sort
s Reft
r) = Sort -> Reft -> SortedReft
RR Sort
s forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc Reft
r

instance Defunc (Symbol, SortedReft) where
  defunc :: (Symbol, SortedReft) -> DF (Symbol, SortedReft)
defunc (Symbol
x, SortedReft
sr) = (Symbol
x,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc SortedReft
sr

instance Defunc (Symbol, Sort) where
  defunc :: (Symbol, Sort) -> DF (Symbol, Sort)
defunc (Symbol
x, Sort
t) = (Symbol
x,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc Sort
t

instance Defunc Reft where
  defunc :: Reft -> DF Reft
defunc (Reft (Symbol
x, Expr
e)) = (Symbol, Expr) -> Reft
Reft forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol
x,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc Expr
e

instance Defunc Expr where
  defunc :: Expr -> DF Expr
defunc = Expr -> DF Expr
txExpr

instance Defunc a => Defunc (SEnv a) where
  defunc :: SEnv a -> DF (SEnv a)
defunc = forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> SEnv a -> m (SEnv b)
mapMSEnv forall a. Defunc a => a -> DF a
defunc

instance Defunc (BindEnv a) where
  defunc :: BindEnv a -> DF (BindEnv a)
defunc BindEnv a
bs = do IBindEnv
dfbs <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets DFST -> IBindEnv
dfBEnv
                 let f :: (Int, (t, SortedReft))
-> StateT DFST Identity (Int, (t, SortedReft))
f (Int
i, (t, SortedReft)
xs) = if Int
i Int -> IBindEnv -> Bool
`memberIBindEnv` IBindEnv
dfbs
                                       then  (Int
i,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc (t, SortedReft)
xs
                                       else  (Int
i,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {t}. (t, SortedReft) -> StateT DFST Identity (t, SortedReft)
matchSort (t, SortedReft)
xs
                 forall (m :: * -> *) a.
Monad m =>
((Int, (Symbol, SortedReft)) -> m (Int, (Symbol, SortedReft)))
-> BindEnv a -> m (BindEnv a)
mapWithKeyMBindEnv forall {t}.
Defunc (t, SortedReft) =>
(Int, (t, SortedReft))
-> StateT DFST Identity (Int, (t, SortedReft))
f BindEnv a
bs
   where
    -- The refinement cannot be elaborated thus defunc-ed because
    -- the bind does not appear in any contraint,
    -- thus unique binders does not perform properly
    -- The sort should be defunc, to ensure same sort on double binders
    matchSort :: (t, SortedReft) -> StateT DFST Identity (t, SortedReft)
matchSort (t
x, RR Sort
s Reft
r) = (t
x,) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Sort -> Reft -> SortedReft
`RR` Reft
r) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Defunc a => a -> DF a
defunc Sort
s

-- Sort defunctionalization [should be done by elaboration]
instance Defunc Sort where
  defunc :: Sort -> DF Sort
defunc = forall (m :: * -> *) a. Monad m => a -> m a
return

instance Defunc a => Defunc [a] where
  defunc :: [a] -> DF [a]
defunc = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall a. Defunc a => a -> DF a
defunc

instance (Defunc a, Eq k, Hashable k) => Defunc (M.HashMap k a) where
  defunc :: HashMap k a -> DF (HashMap k a)
defunc HashMap k a
m = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall (f :: * -> *) b c a.
Functor f =>
(b -> f c) -> (a, b) -> f (a, c)
secondM forall a. Defunc a => a -> DF a
defunc) (forall k v. HashMap k v -> [(k, v)]
M.toList HashMap k a
m)

type DF = State DFST

data DFST = DFST
  { DFST -> Int
dfFresh :: !Int
  , DFST -> SymEnv
dfEnv   :: !SymEnv
  , DFST -> IBindEnv
dfBEnv  :: !IBindEnv
  , DFST -> Bool
dfHO    :: !Bool        -- ^ allow higher order thus defunctionalize
  , DFST -> [Expr]
dfLams  :: ![Expr]      -- ^ lambda expressions appearing in the expressions
  , DFST -> [Expr]
dfRedex :: ![Expr]      -- ^ redexes appearing in the expressions
  , DFST -> SEnv Sort
dfBinds :: !(SEnv Sort) -- ^ sorts of new lambda-binders
  }

makeDFState :: Config -> SymEnv -> IBindEnv -> DFST
makeDFState :: Config -> SymEnv -> IBindEnv -> DFST
makeDFState Config
cfg SymEnv
env IBindEnv
ibind = DFST
  { dfFresh :: Int
dfFresh = Int
0
  , dfEnv :: SymEnv
dfEnv   = SymEnv
env
  , dfBEnv :: IBindEnv
dfBEnv  = IBindEnv
ibind
  , dfHO :: Bool
dfHO    = Config -> Bool
allowHO Config
cfg  Bool -> Bool -> Bool
|| Config -> Bool
defunction Config
cfg
  -- INVARIANT: lambdas and redexes are not defunctionalized
  , dfLams :: [Expr]
dfLams  = []
  , dfRedex :: [Expr]
dfRedex = []
  , dfBinds :: SEnv Sort
dfBinds = forall a. Monoid a => a
mempty
  }

makeInitDFState :: Config -> SInfo a -> DFST
makeInitDFState :: forall a. Config -> SInfo a -> DFST
makeInitDFState Config
cfg SInfo a
si
  = Config -> SymEnv -> IBindEnv -> DFST
makeDFState Config
cfg
      (forall a. Config -> SInfo a -> SymEnv
symbolEnv Config
cfg SInfo a
si)
      (forall a. Monoid a => [a] -> a
mconcat ((forall (c :: * -> *) a. TaggedC c a => c a -> IBindEnv
senv forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. HashMap k v -> [v]
M.elems (forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
Types.cm SInfo a
si)) forall a. [a] -> [a] -> [a]
++ (forall a. WfC a -> IBindEnv
wenv forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. HashMap k v -> [v]
M.elems (forall (c :: * -> *) a. GInfo c a -> HashMap KVar (WfC a)
Types.ws SInfo a
si))))

--------------------------------------------------------------------------------
-- | Low level monad manipulation ----------------------------------------------
--------------------------------------------------------------------------------
freshSym :: Sort -> DF Symbol
freshSym :: Sort -> DF Symbol
freshSym Sort
t = do
  Int
n    <- forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets DFST -> Int
dfFresh
  let x :: Symbol
x = forall a. Show a => Symbol -> a -> Symbol
intSymbol Symbol
"lambda_fun_" Int
n
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify forall a b. (a -> b) -> a -> b
$ \DFST
s -> DFST
s {dfFresh :: Int
dfFresh = Int
n forall a. Num a => a -> a -> a
+ Int
1, dfBinds :: SEnv Sort
dfBinds = forall a. Symbol -> a -> SEnv a -> SEnv a
insertSEnv Symbol
x Sort
t (DFST -> SEnv Sort
dfBinds DFST
s)}
  forall (m :: * -> *) a. Monad m => a -> m a
return Symbol
x