{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE PatternGuards #-}

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

-- This module makes it so no binds with different sorts have the same name.

module Language.Fixpoint.Solver.UniqifyBinds (renameAll) where

import           Language.Fixpoint.Types
import           Language.Fixpoint.Solver.Sanitize (dropDeadSubsts)
import           Language.Fixpoint.Misc          (fst3, mlookup, snd3)

import qualified Data.HashMap.Strict as M
import qualified Data.HashSet        as S
import qualified Data.List           as L
import           Data.Foldable       (foldl')
import           Data.Maybe          (catMaybes, mapMaybe, fromJust, isJust)
import           Data.Hashable       (Hashable)
import           GHC.Generics        (Generic)
import           Control.DeepSeq     (NFData, ($!!))
-- import Debug.Trace (trace)

--------------------------------------------------------------------------------
renameAll    :: (NFData a) => SInfo a -> SInfo a
--------------------------------------------------------------------------------
renameAll :: forall a. NFData a => SInfo a -> SInfo a
renameAll SInfo a
fi2 = SInfo a
fi6
  where
    fi6 :: SInfo a
fi6       = {- SCC "dropDead"    -} forall a. SInfo a -> SInfo a
dropDeadSubsts  SInfo a
fi5
    fi5 :: SInfo a
fi5       = {- SCC "dropUnused"  -} forall a. SInfo a -> SInfo a
dropUnusedBinds SInfo a
fi4
    fi4 :: SInfo a
fi4       = {- SCC "renameBinds" -} forall a. SInfo a -> RenameMap -> SInfo a
renameBinds SInfo a
fi3 forall a b. NFData a => (a -> b) -> a -> b
$!! RenameMap
rnm
    fi3 :: SInfo a
fi3       = {- SCC "renameVars"  -} forall a. SInfo a -> RenameMap -> IdMap -> SInfo a
renameVars SInfo a
fi2 RenameMap
rnm forall a b. NFData a => (a -> b) -> a -> b
$!! IdMap
idm
    rnm :: RenameMap
rnm       = {- SCC "mkRenameMap" -} forall a. BindEnv a -> RenameMap
mkRenameMap forall a b. NFData a => (a -> b) -> a -> b
$!! forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi2
    idm :: IdMap
idm       = {- SCC "mkIdMap"     -} forall a. SInfo a -> IdMap
mkIdMap SInfo a
fi2


--------------------------------------------------------------------------------
-- | `dropUnusedBinds` replaces the refinements of "unused" binders with "true".
--   see tests/pos/unused.fq for an example of why this phase is needed.
--------------------------------------------------------------------------------
dropUnusedBinds :: SInfo a -> SInfo a
dropUnusedBinds :: forall a. SInfo a -> SInfo a
dropUnusedBinds SInfo a
fi = SInfo a
fi {bs :: BindEnv a
bs = forall a.
(BindId -> Symbol -> SortedReft -> Bool) -> BindEnv a -> BindEnv a
filterBindEnv forall {r} {p}. Reftable r => BindId -> p -> r -> Bool
isUsed (forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi)}-- { bs = mapBindEnv tx (bs fi) }
  where
    -- _tx i (x, r)
    -- | isUsed i    = (x, r)
    -- | otherwise   = (x, top r)
    isUsed :: BindId -> p -> r -> Bool
isUsed BindId
i p
_x r
r  = {- tracepp (unwords ["isUsed", show i, showpp _x]) $ -} BindId -> IBindEnv -> Bool
memberIBindEnv BindId
i IBindEnv
usedBinds Bool -> Bool -> Bool
|| forall r. Reftable r => r -> Bool
isTauto r
r
    usedBinds :: IBindEnv
usedBinds      = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' IBindEnv -> IBindEnv -> IBindEnv
unionIBindEnv IBindEnv
emptyIBindEnv ([IBindEnv]
cEnvs forall a. [a] -> [a] -> [a]
++ [IBindEnv]
wEnvs)
    wEnvs :: [IBindEnv]
wEnvs          = 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)
ws SInfo a
fi)
    cEnvs :: [IBindEnv]
cEnvs          = 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 Integer (c a)
cm SInfo a
fi)

data Ref
  = RB !BindId    -- ^ Bind identifier
  | RI !Integer   -- ^ Constraint identifier?
    deriving (Ref -> Ref -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Ref -> Ref -> Bool
$c/= :: Ref -> Ref -> Bool
== :: Ref -> Ref -> Bool
$c== :: Ref -> Ref -> Bool
Eq, forall x. Rep Ref x -> Ref
forall x. Ref -> Rep Ref x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep Ref x -> Ref
$cfrom :: forall x. Ref -> Rep Ref x
Generic)

instance NFData   Ref
instance Hashable Ref

-- | An `IdMap` stores for each constraint and BindId the
--   set of other BindIds that it references, i.e. those
--   where it needs to know when their names gets changed
type IdMap = M.HashMap Ref (S.HashSet BindId)

-- | A `RenameMap` maps an old name and sort to new name,
--   represented by a hashmap containing association lists.
--   `Nothing` as new name means the name is the same as the old.
type RenameMap = M.HashMap Symbol [(Sort, Maybe Symbol)]

--------------------------------------------------------------------------------
mkIdMap :: SInfo a -> IdMap
--------------------------------------------------------------------------------
mkIdMap :: forall a. SInfo a -> IdMap
mkIdMap SInfo a
fi = forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
M.foldlWithKey' (forall a. BindEnv a -> IdMap -> Integer -> SimpC a -> IdMap
updateIdMap forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi) forall k v. HashMap k v
M.empty forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. GInfo c a -> HashMap Integer (c a)
cm SInfo a
fi

updateIdMap :: BindEnv a -> IdMap -> Integer -> SimpC a -> IdMap
updateIdMap :: forall a. BindEnv a -> IdMap -> Integer -> SimpC a -> IdMap
updateIdMap BindEnv a
be IdMap
m Integer
scId SimpC a
s = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
M.insertWith forall a. (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a
S.union (Integer -> Ref
RI Integer
scId) HashSet BindId
refSet IdMap
m'
  where
    ids :: [BindId]
ids                 = IBindEnv -> [BindId]
elemsIBindEnv (forall (c :: * -> *) a. TaggedC c a => c a -> IBindEnv
senv SimpC a
s)
    nameMap :: HashMap Symbol BindId
nameMap             = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList [(forall a b c. (a, b, c) -> a
fst3 forall a b. (a -> b) -> a -> b
$ forall a. BindId -> BindEnv a -> (Symbol, SortedReft, a)
lookupBindEnv BindId
i BindEnv a
be, BindId
i) | BindId
i <- [BindId]
ids]
    m' :: IdMap
m'                  = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a.
BindEnv a -> HashMap Symbol BindId -> IdMap -> BindId -> IdMap
insertIdIdLinks BindEnv a
be HashMap Symbol BindId
nameMap) IdMap
m [BindId]
ids
    symSet :: HashSet Symbol
symSet              = forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a. Subable a => a -> [Symbol]
syms forall a b. (a -> b) -> a -> b
$ forall (c :: * -> *) a. TaggedC c a => c a -> Expr
crhs SimpC a
s
    refSet :: HashSet BindId
refSet              = HashSet Symbol -> HashMap Symbol BindId -> HashSet BindId
namesToIds HashSet Symbol
symSet HashMap Symbol BindId
nameMap

insertIdIdLinks :: BindEnv a -> M.HashMap Symbol BindId -> IdMap -> BindId -> IdMap
insertIdIdLinks :: forall a.
BindEnv a -> HashMap Symbol BindId -> IdMap -> BindId -> IdMap
insertIdIdLinks BindEnv a
be HashMap Symbol BindId
nameMap IdMap
m BindId
i = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
M.insertWith forall a. (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a
S.union (BindId -> Ref
RB BindId
i) HashSet BindId
refSet IdMap
m
  where
    sr :: SortedReft
sr     = forall a b c. (a, b, c) -> b
snd3 forall a b. (a -> b) -> a -> b
$ forall a. BindId -> BindEnv a -> (Symbol, SortedReft, a)
lookupBindEnv BindId
i BindEnv a
be
    symSet :: HashSet Symbol
symSet = Reft -> HashSet Symbol
reftFreeVars forall a b. (a -> b) -> a -> b
$ SortedReft -> Reft
sr_reft SortedReft
sr
    refSet :: HashSet BindId
refSet = HashSet Symbol -> HashMap Symbol BindId -> HashSet BindId
namesToIds HashSet Symbol
symSet HashMap Symbol BindId
nameMap

namesToIds :: S.HashSet Symbol -> M.HashMap Symbol BindId -> S.HashSet BindId
namesToIds :: HashSet Symbol -> HashMap Symbol BindId -> HashSet BindId
namesToIds HashSet Symbol
xs HashMap Symbol BindId
m = forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a. [Maybe a] -> [a]
catMaybes [forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Symbol
x HashMap Symbol BindId
m | Symbol
x <- forall a. HashSet a -> [a]
S.toList HashSet Symbol
xs] --TODO why any Nothings?

--------------------------------------------------------------------------------
mkRenameMap :: BindEnv a -> RenameMap
--------------------------------------------------------------------------------
mkRenameMap :: forall a. BindEnv a -> RenameMap
mkRenameMap BindEnv a
be = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a. BindEnv a -> RenameMap -> BindId -> RenameMap
addId BindEnv a
be) forall k v. HashMap k v
M.empty [BindId]
ids
  where
    ids :: [BindId]
ids = forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. BindEnv a -> [(BindId, (Symbol, SortedReft, a))]
bindEnvToList BindEnv a
be

addId :: BindEnv a -> RenameMap -> BindId -> RenameMap
addId :: forall a. BindEnv a -> RenameMap -> BindId -> RenameMap
addId BindEnv a
be RenameMap
m BindId
i
  | forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member Symbol
sym RenameMap
m = RenameMap -> Symbol -> Sort -> BindId -> RenameMap
addDupId RenameMap
m Symbol
sym Sort
t BindId
i
  | Bool
otherwise      = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert Symbol
sym [(Sort
t, forall a. Maybe a
Nothing)] RenameMap
m
  where
    t :: Sort
t              = SortedReft -> Sort
sr_sort SortedReft
sr
    (Symbol
sym, SortedReft
sr, a
_)   = forall a. BindId -> BindEnv a -> (Symbol, SortedReft, a)
lookupBindEnv BindId
i BindEnv a
be

addDupId :: RenameMap -> Symbol -> Sort -> BindId -> RenameMap
addDupId :: RenameMap -> Symbol -> Sort -> BindId -> RenameMap
addDupId RenameMap
m Symbol
sym Sort
t BindId
i
  | forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup Sort
t [(Sort, Maybe Symbol)]
mapping = RenameMap
m
  | Bool
otherwise                   = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert Symbol
sym ((Sort
t, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Symbol -> BindId -> Symbol
renameSymbol Symbol
sym BindId
i) forall a. a -> [a] -> [a]
: [(Sort, Maybe Symbol)]
mapping) RenameMap
m
  where
    mapping :: [(Sort, Maybe Symbol)]
mapping = forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Symbol
sym RenameMap
m

--------------------------------------------------------------------------------
-- | `renameVars` seems to do the actual work of renaming all the binders
--   to use their sort-specific names.
--------------------------------------------------------------------------------
renameVars :: SInfo a -> RenameMap -> IdMap -> SInfo a
--------------------------------------------------------------------------------
renameVars :: forall a. SInfo a -> RenameMap -> IdMap -> SInfo a
renameVars SInfo a
fi RenameMap
rnMap IdMap
idMap = forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
M.foldlWithKey' (forall a. RenameMap -> SInfo a -> Ref -> HashSet BindId -> SInfo a
updateRef RenameMap
rnMap) SInfo a
fi IdMap
idMap

updateRef :: RenameMap -> SInfo a -> Ref -> S.HashSet BindId -> SInfo a
updateRef :: forall a. RenameMap -> SInfo a -> Ref -> HashSet BindId -> SInfo a
updateRef RenameMap
rnMap SInfo a
fi Ref
rf HashSet BindId
bset = forall a. Subst -> SInfo a -> Ref -> SInfo a
applySub ([(Symbol, Expr)] -> Subst
mkSubst [(Symbol, Expr)]
subs) SInfo a
fi Ref
rf
  where
    symTList :: [(Symbol, Sort)]
symTList = [ (Symbol
sym, SortedReft -> Sort
sr_sort SortedReft
sr) | BindId
i <- forall a. HashSet a -> [a]
S.toList HashSet BindId
bset, let (Symbol
sym, SortedReft
sr, a
_) = forall a. BindId -> BindEnv a -> (Symbol, SortedReft, a)
lookupBindEnv BindId
i BindEnv a
bEnv]
    bEnv :: BindEnv a
bEnv     = forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi
    subs :: [(Symbol, Expr)]
subs     = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (RenameMap -> (Symbol, Sort) -> Maybe (Symbol, Expr)
mkSubUsing RenameMap
rnMap) [(Symbol, Sort)]
symTList

mkSubUsing :: RenameMap -> (Symbol, Sort) -> Maybe (Symbol, Expr)
mkSubUsing :: RenameMap -> (Symbol, Sort) -> Maybe (Symbol, Expr)
mkSubUsing RenameMap
m (Symbol
sym, Sort
t) = do
  Symbol
newName <- forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup Sort
t forall a b. (a -> b) -> a -> b
$ forall k v.
(HasCallStack, Eq k, Show k, Hashable k) =>
HashMap k v -> k -> v
mlookup RenameMap
m Symbol
sym
  forall (m :: * -> *) a. Monad m => a -> m a
return (Symbol
sym, forall a. Symbolic a => a -> Expr
eVar Symbol
newName)

applySub :: Subst -> SInfo a -> Ref -> SInfo a
applySub :: forall a. Subst -> SInfo a -> Ref -> SInfo a
applySub Subst
sub SInfo a
fi (RB BindId
i) = SInfo a
fi { bs :: BindEnv a
bs = forall a.
((Symbol, SortedReft) -> (Symbol, SortedReft))
-> BindId -> BindEnv a -> BindEnv a
adjustBindEnv forall {b} {a}. Subable b => (a, b) -> (a, b)
go BindId
i (forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi) }
  where
    go :: (a, b) -> (a, b)
go (a
sym, b
sr)       = (a
sym, forall a. Subable a => Subst -> a -> a
subst Subst
sub b
sr)

applySub Subst
sub SInfo a
fi (RI Integer
i) = SInfo a
fi { cm :: HashMap Integer (SimpC a)
cm = forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
M.adjust forall {a}. SimpC a -> SimpC a
go Integer
i (forall (c :: * -> *) a. GInfo c a -> HashMap Integer (c a)
cm SInfo a
fi) }
  where
    go :: SimpC a -> SimpC a
go SimpC a
c               = SimpC a
c { _crhs :: Expr
_crhs = forall a. Subable a => Subst -> a -> a
subst Subst
sub (forall a. SimpC a -> Expr
_crhs SimpC a
c) }

--------------------------------------------------------------------------------
renameBinds :: SInfo a -> RenameMap -> SInfo a
--------------------------------------------------------------------------------
renameBinds :: forall a. SInfo a -> RenameMap -> SInfo a
renameBinds SInfo a
fi RenameMap
m = SInfo a
fi { bs :: BindEnv a
bs = forall a. [(BindId, (Symbol, SortedReft, a))] -> BindEnv a
bindEnvFromList forall a b. (a -> b) -> a -> b
$ forall a.
RenameMap
-> (BindId, (Symbol, SortedReft, a))
-> (BindId, (Symbol, SortedReft, a))
renameBind RenameMap
m forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(BindId, (Symbol, SortedReft, a))]
beList }
  where
    beList :: [(BindId, (Symbol, SortedReft, a))]
beList       = forall a. BindEnv a -> [(BindId, (Symbol, SortedReft, a))]
bindEnvToList (forall (c :: * -> *) a. GInfo c a -> BindEnv a
bs SInfo a
fi)

renameBind :: RenameMap -> (BindId, (Symbol, SortedReft, a)) -> (BindId, (Symbol, SortedReft, a))
renameBind :: forall a.
RenameMap
-> (BindId, (Symbol, SortedReft, a))
-> (BindId, (Symbol, SortedReft, a))
renameBind RenameMap
m (BindId
i, (Symbol
sym, SortedReft
sr, a
ann))
  | Just Symbol
newSym <- Maybe Symbol
mnewSym = (BindId
i, (Symbol
newSym, SortedReft
sr, a
ann))
  | Bool
otherwise              = (BindId
i, (Symbol
sym,    SortedReft
sr, a
ann))
  where
    t :: Sort
t                      = SortedReft -> Sort
sr_sort SortedReft
sr
    mnewSym :: Maybe Symbol
mnewSym                = forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall a b. Eq a => a -> [(a, b)] -> Maybe b
L.lookup Sort
t forall a b. (a -> b) -> a -> b
$ forall k v.
(HasCallStack, Eq k, Show k, Hashable k) =>
HashMap k v -> k -> v
mlookup RenameMap
m Symbol
sym