-- | 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 ()

-- | 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 = Var -> TagEnv -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member

makeTagEnv :: [CoreBind] -> TagEnv
makeTagEnv :: [CoreBind] -> TagEnv
makeTagEnv = {- tracepp "TAGENV" . -} (Int -> Tag) -> HashMap Var Int -> TagEnv
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.map (Int -> Tag -> Tag
forall a. a -> [a] -> [a]
:[]) (HashMap Var Int -> TagEnv)
-> ([CoreBind] -> HashMap Var Int) -> [CoreBind] -> TagEnv
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> HashMap Var Int
callGraphRanks (CallGraph -> HashMap Var Int)
-> ([CoreBind] -> CallGraph) -> [CoreBind] -> HashMap Var Int
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 = Tag -> Var -> TagEnv -> Tag
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  = [(Var, Int)] -> HashMap Var Int
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList ([(Var, Int)] -> HashMap Var Int)
-> (CallGraph -> [(Var, Int)]) -> CallGraph -> HashMap Var Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[(Var, Int)]] -> [(Var, Int)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Var, Int)]] -> [(Var, Int)])
-> (CallGraph -> [[(Var, Int)]]) -> CallGraph -> [(Var, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [SCC Var] -> [[(Var, Int)]]
forall {a}. [SCC a] -> [[(a, Int)]]
index ([SCC Var] -> [[(Var, Int)]])
-> (CallGraph -> [SCC Var]) -> CallGraph -> [[(Var, Int)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CallGraph -> [SCC Var]
forall {key}. Ord key => [(key, [key])] -> [SCC key]
mkScc
  where mkScc :: [(key, [key])] -> [SCC key]
mkScc [(key, [key])]
cg = [(key, key, [key])] -> [SCC key]
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    = (Int -> SCC a -> [(a, Int)]) -> Tag -> [SCC a] -> [[(a, Int)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i -> (a -> (a, Int)) -> [a] -> [(a, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (, Int
i) ([a] -> [(a, Int)]) -> (SCC a -> [a]) -> SCC a -> [(a, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SCC a -> [a]
forall vertex. SCC vertex -> [vertex]
G.flattenSCC) [Int
1..]

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