g2-0.2.0.0: src/G2/Liquid/Inference/PolyRef.hs
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable#-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TupleSections #-}
module G2.Liquid.Inference.PolyRef ( PolyBound (.. )
, RefNamePolyBound
, ExprPolyBound
, extractExprPolyBoundWithRoot
, extractExprPolyBound
, extractTypePolyBound
, extractTypeAppAndFuncPolyBound
, headValue
, removeHead
, extractValues
, uniqueIds
, mapPB
, filterPB
, allPB
, zipPB
, zipWithPB
, zipWithMaybePB
, zip3PB) where
import G2.Language
import qualified Data.HashMap.Lazy as HM
import Data.List
import Data.Maybe
type RefNamePolyBound = PolyBound String
type ExprPolyBound = PolyBound [Expr]
type TypePolyBound = PolyBound Type
-- | The subexpressions of an expression corresponding to the polymorphic
-- arguments. If a polymorphic argument is instantiated with a polymorphic
-- type, these are nested recursively.
data PolyBound v = PolyBound v [PolyBound v] deriving (Eq, Read, Show, Functor, Foldable, Traversable)
-------------------------------
-- ExprPolyBound
-------------------------------
extractExprPolyBoundWithRoot :: Expr -> ExprPolyBound
extractExprPolyBoundWithRoot e = PolyBound [e] $ extractExprPolyBound e
extractExprPolyBound :: Expr -> [ExprPolyBound]
extractExprPolyBound e
| Data dc:_ <- unApp e =
let
bound = leadingTyForAllBindings dc
m = extractExprPolyBound' e
bound_es = map (\i -> HM.lookupDefault [] i m) bound
in
map (\es -> PolyBound es (mergeExprPolyBound . transpose $ map extractExprPolyBound es)) bound_es
| otherwise = []
mergeExprPolyBound :: [[ExprPolyBound]] -> [ExprPolyBound]
mergeExprPolyBound = mapMaybe (\pb -> case pb of
(p:pbb) -> Just $ foldr mergeExprPolyBound' p pbb
[] -> Nothing)
mergeExprPolyBound' :: ExprPolyBound -> ExprPolyBound -> ExprPolyBound
mergeExprPolyBound' (PolyBound es1 pb1) (PolyBound es2 pb2) =
PolyBound (es1 ++ es2) (map (uncurry mergeExprPolyBound') $ zip pb1 pb2)
extractExprPolyBound' :: Expr -> HM.HashMap Id [Expr]
extractExprPolyBound' e
| Data dc:es <- unApp e =
let
es' = filter (not . isType) es
argtys = argumentTypes . PresType . inTyForAlls $ typeOf dc
argtys_es = zip argtys es'
(direct, indirect) = partition fstIsTyVar argtys_es
direct' = mapMaybe fstMapTyVar direct
indirect' = map (uncurry substTypes) indirect
direct_hm = foldr (HM.unionWith (++)) HM.empty
$ map (\(i, e_) -> uncurry HM.singleton (i, e_:[])) direct'
in
foldr (HM.unionWith (++)) direct_hm $ map (extractExprPolyBound' . adjustIndirectTypes) indirect'
| otherwise = HM.empty
where
isType (Type _) = True
isType _ = False
fstIsTyVar (TyVar _, _) = True
fstIsTyVar _ = False
fstMapTyVar (TyVar i, x) = Just (i, x)
fstMapTyVar _ = Nothing
substTypes :: Type -> Expr -> Expr
substTypes t e
| _:ts <- unTyApp t
, e':es <- unApp e =
mkApp $ e':substTypes' ts es
substTypes _ e = e
substTypes' :: [Type] -> [Expr] -> [Expr]
substTypes' (t:ts) (Type _:es) = Type t:substTypes' ts es
substTypes' _ es = es
adjustIndirectTypes :: Expr -> Expr
adjustIndirectTypes e
| Data dc:es <- unApp e =
let
tyses = filter (isType) es
tyses' = map (\(Type t) -> t) tyses
bound = leadingTyForAllBindings dc
bound_tyses = zip bound tyses'
in
mkApp $ Data (foldr (uncurry retype) dc $ bound_tyses):es
| otherwise = e
where
isType (Type _) = True
isType _ = False
-------------------------------
-- TypePolyBound
-------------------------------
-- | Unrolls TyApp'ed args, while also keeping them in the base type
extractTypePolyBound :: Type -> TypePolyBound
extractTypePolyBound t =
let
(_:ts) = unTyApp t
in
PolyBound t $ map extractTypePolyBound ts
-- | Unrolls TyApp'ed and TyFunc'ed args, while also keeping them in the base type
extractTypeAppAndFuncPolyBound :: Type -> TypePolyBound
extractTypeAppAndFuncPolyBound t@(TyApp _ _) =
let
(_:ts) = unTyApp t
in
PolyBound t $ map extractTypePolyBound ts
extractTypeAppAndFuncPolyBound t@(TyFun _ _) =
let
ts = splitTyFuns t
in
PolyBound t $ map extractTypePolyBound ts
extractTypeAppAndFuncPolyBound t = PolyBound t []
-------------------------------
-- Generic PolyBound functions
-------------------------------
headValue :: PolyBound v -> v
headValue (PolyBound v _) = v
removeHead :: PolyBound v -> [PolyBound v]
removeHead (PolyBound _ vs) = vs
extractValues :: PolyBound v -> [v]
extractValues (PolyBound v ps) = v:concatMap extractValues ps
uniqueIds :: PolyBound v -> PolyBound Int
uniqueIds = snd . uniqueIds' 0
uniqueIds' :: Int -> PolyBound v -> (Int, PolyBound Int)
uniqueIds' n (PolyBound _ ps) =
let
(n', ps') = mapAccumR (uniqueIds') (n + 1) ps
in
(n', PolyBound n ps')
mapPB :: (a -> b) -> PolyBound a -> PolyBound b
mapPB f (PolyBound v ps) = PolyBound (f v) (map (mapPB f) ps)
filterPB :: (PolyBound a -> Bool) -> PolyBound a -> Maybe (PolyBound a)
filterPB p pb@(PolyBound v xs) =
case p pb of
True -> Just $ PolyBound v (mapMaybe (filterPB p) xs)
False -> Nothing
allPB :: (a -> Bool) -> PolyBound a -> Bool
allPB p = all p . extractValues
zipPB :: PolyBound a -> PolyBound b -> PolyBound (a, b)
zipPB (PolyBound a pba) (PolyBound b pbb) = PolyBound (a, b) (zipWith zipPB pba pbb)
zipWithPB :: (a -> b -> c) -> PolyBound a -> PolyBound b -> PolyBound c
zipWithPB f (PolyBound a pba) (PolyBound b pbb) = PolyBound (f a b) (zipWith (zipWithPB f) pba pbb)
zipWithMaybePB :: (Maybe a -> Maybe b -> c) -> PolyBound a -> PolyBound b -> PolyBound c
zipWithMaybePB f pba pbb = zipWithMaybePB' f (mapPB Just pba) (mapPB Just pbb)
zipWithMaybePB' :: (Maybe a -> Maybe b -> c) -> PolyBound (Maybe a) -> PolyBound (Maybe b) -> PolyBound c
zipWithMaybePB' f (PolyBound a pba) (PolyBound b pbb) =
let
c = f a b
rep_nt = repeat (PolyBound Nothing [])
pbc = takeWhile (\(x, y) -> isJust (headValue x) || isJust (headValue y))
$ zip (pba ++ rep_nt) (pbb ++ rep_nt)
in
PolyBound c $ map (uncurry (zipWithMaybePB' f)) pbc
zip3PB :: PolyBound a -> PolyBound b -> PolyBound c -> PolyBound (a, b, c)
zip3PB (PolyBound a pba) (PolyBound b pbb) (PolyBound c pbc) =
PolyBound (a, b, c) (zipWith3 zip3PB pba pbb pbc)