-- | This module contains the code for generating "tags" for constraints
--   based on their source, i.e. the top-level binders under which the
--   constraint was generated. These tags are used by fixpoint to
--   prioritize constraints by the "source-level" function.

{-# LANGUAGE TupleSections #-}

module Language.Haskell.Liquid.UX.CTags (
    -- * Type for constraint tags
    TagKey, TagEnv

    -- * Default tag value
  , defaultTag

    -- * Constructing @TagEnv@
  , makeTagEnv

    -- * Accessing @TagEnv@
  , getTag
  , memTagEnv

) where

import Prelude hiding (error)

import qualified Data.HashSet           as S
import qualified Data.HashMap.Strict    as M
import qualified Data.Graph             as G

import Language.Fixpoint.Types          (Tag)
import Liquid.GHC.API
import Language.Haskell.Liquid.Types.Visitors (freeVars)
import Language.Haskell.Liquid.Types.PrettyPrint ()
import Language.Fixpoint.Misc     (mapSnd)

-- | The @TagKey@ is the top-level binder, and @Tag@ is a singleton Int list

type TagKey = Var
type TagEnv = M.HashMap TagKey Tag

-- TODO: use the "callgraph" SCC to do this numbering.

defaultTag :: Tag
defaultTag :: Tag
defaultTag = [Int
0]

memTagEnv :: TagKey -> TagEnv -> Bool
memTagEnv :: Var -> TagEnv -> Bool
memTagEnv = forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member

makeTagEnv :: [CoreBind] -> TagEnv
makeTagEnv :: [CoreBind] -> TagEnv
makeTagEnv = {- tracepp "TAGENV" . -} forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.map (forall a. a -> [a] -> [a]
:[]) forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> HashMap Var Int
callGraphRanks forall b c a. (b -> c) -> (a -> b) -> a -> c
. [CoreBind] -> CallGraph
makeCallGraph

-- makeTagEnv = M.fromList . (`zip` (map (:[]) [1..])). L.sort . map fst . concatMap bindEqns

getTag :: TagKey -> TagEnv -> Tag
getTag :: Var -> TagEnv -> Tag
getTag = forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
M.lookupDefault Tag
defaultTag

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

type CallGraph = [(Var, [Var])] -- caller-callee pairs

callGraphRanks :: CallGraph -> M.HashMap Var Int
-- callGraphRanks cg = traceShow ("CallGraph Ranks: " ++ show cg) $ callGraphRanks' cg

callGraphRanks :: CallGraph -> HashMap Var Int
callGraphRanks  = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. [SCC a] -> [[(a, Int)]]
index forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {key}. Ord key => [(key, [key])] -> [SCC key]
mkScc
  where mkScc :: [(key, [key])] -> [SCC key]
mkScc [(key, [key])]
cg = forall key node. Ord key => [(node, key, [key])] -> [SCC node]
G.stronglyConnComp [(key
u, key
u, [key]
vs) | (key
u, [key]
vs) <- [(key, [key])]
cg]
        index :: [SCC a] -> [[(a, Int)]]
index    = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i -> forall a b. (a -> b) -> [a] -> [b]
map (, Int
i) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall vertex. SCC vertex -> [vertex]
G.flattenSCC) [Int
1..]

makeCallGraph :: [CoreBind] -> CallGraph
makeCallGraph :: [CoreBind] -> CallGraph
makeCallGraph [CoreBind]
cbs = forall b c a. (b -> c) -> (a, b) -> (a, c)
mapSnd Expr Var -> [Var]
calls forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` [(Var, Expr Var)]
xes
  where xes :: [(Var, Expr Var)]
xes       = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall t. Bind t -> [(t, Expr t)]
bindEqns [CoreBind]
cbs
        xs :: HashSet Var
xs        = forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Var, Expr Var)]
xes
        calls :: Expr Var -> [Var]
calls     = forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`S.member` HashSet Var
xs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars forall a. HashSet a
S.empty

bindEqns :: Bind t -> [(t, Expr t)]
bindEqns :: forall t. Bind t -> [(t, Expr t)]
bindEqns (NonRec t
x Expr t
e) = [(t
x, Expr t
e)]
bindEqns (Rec [(t, Expr t)]
xes)    = [(t, Expr t)]
xes