{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE TupleSections         #-}

module Gradual.Concretize (Gradual(..)) where

import Gradual.Types 
import Language.Fixpoint.Types
import Gradual.Misc

import qualified Data.HashMap.Strict       as M
-- import Debug.Trace 

class Gradual a where
  concretize :: GMap GWInfo -> a -> [(GSub GWInfo, a)]
  concretize GMap GWInfo
_ a
x = [(GSub GWInfo
forall a. Monoid a => a
mempty, a
x)]

instance Gradual (SInfo a) where
  concretize :: GMap GWInfo -> SInfo a -> [(GSub GWInfo, SInfo a)]
concretize GMap GWInfo
i SInfo a
sinfo = (\GSub GWInfo
su -> (GSub GWInfo
su,
     SInfo a
sinfo {bs :: BindEnv
bs = (BindEnv, GSub GWInfo) -> BindEnv -> BindEnv
forall a. GSubable a => (BindEnv, GSub GWInfo) -> a -> a
gsubst (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
sinfo,GSub GWInfo
su) (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
sinfo), cm :: HashMap SubcId (SimpC a)
cm = (BindEnv, GSub GWInfo)
-> HashMap SubcId (SimpC a) -> HashMap SubcId (SimpC a)
forall a. GSubable a => (BindEnv, GSub GWInfo) -> a -> a
gsubst (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
sinfo,GSub GWInfo
su) (SInfo a -> HashMap SubcId (SimpC a)
forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
cm SInfo a
sinfo)} 
    )) (GSub GWInfo -> (GSub GWInfo, SInfo a))
-> [GSub GWInfo] -> [(GSub GWInfo, SInfo a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([(KVar, (GWInfo, Expr))] -> GSub GWInfo
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList ([(KVar, (GWInfo, Expr))] -> GSub GWInfo)
-> [[(KVar, (GWInfo, Expr))]] -> [GSub GWInfo]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(KVar, (GWInfo, [Expr]))] -> [[(KVar, (GWInfo, Expr))]]
forall k i v. [(k, (i, [v]))] -> [[(k, (i, v))]]
flatten (GMap GWInfo -> [(KVar, (GWInfo, [Expr]))]
forall k v. HashMap k v -> [(k, v)]
M.toList GMap GWInfo
i))


class GSubable a where
  gsubst :: (BindEnv, GSub GWInfo) -> a -> a 
  gsubst (BindEnv, GSub GWInfo)
_ a
x = a
x 

instance GSubable BindEnv where
  gsubst :: (BindEnv, GSub GWInfo) -> BindEnv -> BindEnv
gsubst (BindEnv, GSub GWInfo)
i BindEnv
benv = [(BindId, Symbol, SortedReft)] -> BindEnv
bindEnvFromList ((SortedReft -> SortedReft)
-> (BindId, Symbol, SortedReft) -> (BindId, Symbol, SortedReft)
forall c d a b. (c -> d) -> (a, b, c) -> (a, b, d)
mapThd3 ((BindEnv, GSub GWInfo) -> SortedReft -> SortedReft
forall a. GSubable a => (BindEnv, GSub GWInfo) -> a -> a
gsubst (BindEnv, GSub GWInfo)
i) ((BindId, Symbol, SortedReft) -> (BindId, Symbol, SortedReft))
-> [(BindId, Symbol, SortedReft)] -> [(BindId, Symbol, SortedReft)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> BindEnv -> [(BindId, Symbol, SortedReft)]
bindEnvToList BindEnv
benv)

instance GSubable (SimpC a) where
  gsubst :: (BindEnv, GSub GWInfo) -> SimpC a -> SimpC a
gsubst (BindEnv
benv,GSub GWInfo
i) SimpC a
c = SimpC a
c {_crhs :: Expr
_crhs = Symbol -> GSub GWInfo -> Expr -> Expr
substGrad Symbol
x GSub GWInfo
i (SimpC a -> Expr
forall a. SimpC a -> Expr
_crhs SimpC a
c)}
    where
      x :: Symbol
x = (Symbol, SortedReft) -> Symbol
forall a b. (a, b) -> a
fst ((Symbol, SortedReft) -> Symbol) -> (Symbol, SortedReft) -> Symbol
forall a b. (a -> b) -> a -> b
$ BindId -> BindEnv -> (Symbol, SortedReft)
lookupBindEnv (SimpC a -> BindId
forall a. SimpC a -> BindId
cbind SimpC a
c) BindEnv
benv

instance (GSubable v) => GSubable (M.HashMap SubcId v) where
  gsubst :: (BindEnv, GSub GWInfo) -> HashMap SubcId v -> HashMap SubcId v
gsubst (BindEnv, GSub GWInfo)
i HashMap SubcId v
m = (v -> v) -> HashMap SubcId v -> HashMap SubcId v
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.map ((BindEnv, GSub GWInfo) -> v -> v
forall a. GSubable a => (BindEnv, GSub GWInfo) -> a -> a
gsubst (BindEnv, GSub GWInfo)
i) HashMap SubcId v
m

instance GSubable SortedReft where
  gsubst :: (BindEnv, GSub GWInfo) -> SortedReft -> SortedReft
gsubst (BindEnv, GSub GWInfo)
i (RR Sort
s Reft
r) = Sort -> Reft -> SortedReft
RR Sort
s (Reft -> SortedReft) -> Reft -> SortedReft
forall a b. (a -> b) -> a -> b
$ (BindEnv, GSub GWInfo) -> Reft -> Reft
forall a. GSubable a => (BindEnv, GSub GWInfo) -> a -> a
gsubst (BindEnv, GSub GWInfo)
i Reft
r 

instance GSubable Reft where
  gsubst :: (BindEnv, GSub GWInfo) -> Reft -> Reft
gsubst (BindEnv
_,GSub GWInfo
i) (Reft (Symbol
x,Expr
p)) = (Symbol, Expr) -> Reft
Reft ((Symbol, Expr) -> Reft)
-> (Expr -> (Symbol, Expr)) -> Expr -> Reft
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol
x,) (Expr -> Reft) -> Expr -> Reft
forall a b. (a -> b) -> a -> b
$ Symbol -> GSub GWInfo -> Expr -> Expr
substGrad Symbol
x GSub GWInfo
i Expr
p

-- instance GSubable Expr where
--   gsubst i p = substGrad dummySymbol i p

substGrad :: Symbol -> GSub GWInfo -> Expr -> Expr
substGrad :: Symbol -> GSub GWInfo -> Expr -> Expr
substGrad Symbol
x GSub GWInfo
m (PGrad KVar
k Subst
su GradInfo
_ Expr
e) 
  = case KVar -> GSub GWInfo -> Maybe (GWInfo, Expr)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup KVar
k GSub GWInfo
m of 
      Just (GWInfo
i, Expr
ek) -> [Expr] -> Expr
pAnd [Subst -> Expr -> Expr
forall a. Subable a => Subst -> a -> a
subst Subst
su (Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
$ Expr -> (Symbol, Expr) -> Expr
forall a. Subable a => a -> (Symbol, Expr) -> a
subst1 Expr
ek (GWInfo -> Symbol
gsym GWInfo
i, Symbol -> Expr
EVar Symbol
x), Expr
e] 
                      -- in trace ("SUBSTITUTED: STATIC = " ++ showFix e ++ 
                      --           "\t GRADUAL = " ++ showFix ee ++
                      --           "\t INIT = " ++ showFix ek) ee
      Maybe (GWInfo, Expr)
Nothing      -> Expr
PTrue  
substGrad Symbol
x GSub GWInfo
m (PAnd [Expr]
es)        
  = [Expr] -> Expr
PAnd (Symbol -> GSub GWInfo -> Expr -> Expr
substGrad Symbol
x GSub GWInfo
m (Expr -> Expr) -> [Expr] -> [Expr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Expr]
es)
substGrad Symbol
_ GSub GWInfo
_ Expr
e
  = Expr
e 


{- 
instance Gradual i Expr where
  concretize i (PGrad k su _ _) = subst su <$> (snd $ M.lookupDefault err k i) 
    where
      err = (undefined, [PTrue]) -- error ("Gradual Not found: " ++ show k)
  concretize i (PAnd ps) = PAnd <$> expand (concretize i) ps 
  concretize _ p = [p]


instance Gradual i BindEnv where
  concretize i benv = bindEnvFromList <$> expand3 (concretize i) (bindEnvToList benv)

instance Fixpoint a => Gradual i (SimpC a) where
  concretize i c = [c{_crhs = rhs} | rhs <- concretize i (_crhs c)]

instance (Gradual i v) => Gradual i (M.HashMap SubcId v) where
  concretize i m = M.fromList <$> expand2 (concretize i) (M.toList m)

instance Gradual i SortedReft where
  concretize i (RR s r) = RR s <$> concretize i r 

instance Gradual i Reft where
  concretize i (Reft (x,e)) = (Reft . (x,)) <$> concretize i e  
-}