{-|
  Copyright   :  (C) 2012-2016, University of Twente
  License     :  BSD2 (see the file LICENSE)
  Maintainer  :  Christiaan Baaij <christiaan.baaij@gmail.com>

  Smart constructor and destructor functions for CoreHW
-}

{-# LANGUAGE CPP               #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# LANGUAGE ViewPatterns      #-}

module Clash.Core.Util where

import           Control.Concurrent.Supply     (Supply, freshId)
import qualified Control.Lens                  as Lens
import Control.Monad.Trans.Except              (Except, throwE)
import           Data.Coerce                   (coerce)
import qualified Data.IntSet                   as IntSet
import qualified Data.HashSet                  as HashSet
import Data.List
  (foldl', mapAccumR, elemIndices, nub)
import Data.Maybe                              (fromJust, mapMaybe, catMaybes)
import qualified Data.Text                     as T
import           Data.Text.Prettyprint.Doc     (line)
#if !MIN_VERSION_base(4,11,0)
import           Data.Semigroup
#endif

import Clash.Core.DataCon
  (DataCon (MkData), dcType, dcUnivTyVars, dcExtTyVars, dcArgTys)
import Clash.Core.FreeVars
  (termFreeVarsX, tyFVsOfTypes, typeFreeVars)
import Clash.Core.Literal                      (literalType)
import Clash.Core.Name
  (Name (..), OccName, mkUnsafeInternalName, mkUnsafeSystemName)
import Clash.Core.Pretty                       (ppr, showPpr)
import Clash.Core.Subst
  (extendTvSubst, mkSubst, mkTvSubst, substTy, substTyWith,
   substTyInVar, extendTvSubstList)
import Clash.Core.Term
  (LetBinding, Pat (..), PrimInfo (..), Term (..), Alt, WorkInfo (..),
   TickInfo (..), collectArgs)
import Clash.Core.Type
  (Kind, LitTy (..), Type (..), TypeView (..),
   coreView, coreView1, isFunTy, isPolyFunCoreTy, mkFunTy, splitFunTy, tyView,
   undefinedTy, isTypeFamilyApplication)
import Clash.Core.TyCon
  (TyConMap, tyConDataCons)
import Clash.Core.TysPrim                      (typeNatKind)
import Clash.Core.Var
  (Id, TyVar, Var (..), isLocalId, mkLocalId, mkTyVar)
import Clash.Core.VarEnv
  (InScopeSet, VarEnv, emptyVarEnv, extendInScopeSet, extendVarEnv,
   lookupVarEnv, mkInScopeSet, uniqAway, extendInScopeSetList, unionInScope,
   mkVarSet, unionVarSet, unitVarSet, emptyVarSet)
import Clash.Unique
import Clash.Util

-- | Type environment/context
type Gamma = VarEnv Type
-- | Kind environment/context
type Delta = VarEnv Kind

-- | Given the left and right side of an equation, normalize it such that
-- equations of the following forms:
--
--     * 5     ~ n + 2
--     * 5     ~ 2 + n
--     * n + 2 ~ 5
--     * 2 + n ~ 5
--
-- are returned as (5, 2, n)
normalizeAdd
  :: (Type, Type)
  -> Maybe (Integer, Integer, Type)
normalizeAdd :: (Type, Type) -> Maybe (Integer, Integer, Type)
normalizeAdd (a :: Type
a, b :: Type
b) = do
  (n :: Integer
n, rhs :: Type
rhs) <- Type -> Type -> Maybe (Integer, Type)
lhsLit Type
a Type
b
  case Type -> TypeView
tyView Type
rhs of
    TyConApp (TyConName -> OccName
forall a. Name a -> OccName
nameOcc -> OccName
"GHC.TypeNats.+") [left :: Type
left, right :: Type
right] -> do
      (m :: Integer
m, o :: Type
o) <- Type -> Type -> Maybe (Integer, Type)
lhsLit Type
left Type
right
      (Integer, Integer, Type) -> Maybe (Integer, Integer, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer
n, Integer
m, Type
o)
    _ ->
      Maybe (Integer, Integer, Type)
forall a. Maybe a
Nothing
 where
  lhsLit :: Type -> Type -> Maybe (Integer, Type)
lhsLit x :: Type
x                 (LitTy (NumTy n :: Integer
n)) = (Integer, Type) -> Maybe (Integer, Type)
forall a. a -> Maybe a
Just (Integer
n, Type
x)
  lhsLit (LitTy (NumTy n :: Integer
n)) y :: Type
y                 = (Integer, Type) -> Maybe (Integer, Type)
forall a. a -> Maybe a
Just (Integer
n, Type
y)
  lhsLit _                 _                 = Maybe (Integer, Type)
forall a. Maybe a
Nothing

-- | Data type that indicates what kind of solution (if any) was found
data TypeEqSolution
  = Solution (TyVar, Type)
  -- ^ Solution was found. Variable equals some integer.
  | AbsurdSolution
  -- ^ A solution was found, but it involved negative naturals.
  | NoSolution
  -- ^ Given type wasn't an equation, or it was unsolvable.
    deriving (Int -> TypeEqSolution -> ShowS
[TypeEqSolution] -> ShowS
TypeEqSolution -> String
(Int -> TypeEqSolution -> ShowS)
-> (TypeEqSolution -> String)
-> ([TypeEqSolution] -> ShowS)
-> Show TypeEqSolution
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TypeEqSolution] -> ShowS
$cshowList :: [TypeEqSolution] -> ShowS
show :: TypeEqSolution -> String
$cshow :: TypeEqSolution -> String
showsPrec :: Int -> TypeEqSolution -> ShowS
$cshowsPrec :: Int -> TypeEqSolution -> ShowS
Show, TypeEqSolution -> TypeEqSolution -> Bool
(TypeEqSolution -> TypeEqSolution -> Bool)
-> (TypeEqSolution -> TypeEqSolution -> Bool) -> Eq TypeEqSolution
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TypeEqSolution -> TypeEqSolution -> Bool
$c/= :: TypeEqSolution -> TypeEqSolution -> Bool
== :: TypeEqSolution -> TypeEqSolution -> Bool
$c== :: TypeEqSolution -> TypeEqSolution -> Bool
Eq)

catSolutions :: [TypeEqSolution] -> [(TyVar, Type)]
catSolutions :: [TypeEqSolution] -> [(TyVar, Type)]
catSolutions = (TypeEqSolution -> Maybe (TyVar, Type))
-> [TypeEqSolution] -> [(TyVar, Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe TypeEqSolution -> Maybe (TyVar, Type)
getSol
 where
  getSol :: TypeEqSolution -> Maybe (TyVar, Type)
getSol (Solution s :: (TyVar, Type)
s) = (TyVar, Type) -> Maybe (TyVar, Type)
forall a. a -> Maybe a
Just (TyVar, Type)
s
  getSol _ = Maybe (TyVar, Type)
forall a. Maybe a
Nothing

-- | Solve given equations and return all non-absurd solutions
solveNonAbsurds :: TyConMap -> [(Type, Type)] -> [(TyVar, Type)]
solveNonAbsurds :: TyConMap -> [(Type, Type)] -> [(TyVar, Type)]
solveNonAbsurds _tcm :: TyConMap
_tcm [] = []
solveNonAbsurds tcm :: TyConMap
tcm (eq :: (Type, Type)
eq:eqs :: [(Type, Type)]
eqs) =
  [(TyVar, Type)]
solved [(TyVar, Type)] -> [(TyVar, Type)] -> [(TyVar, Type)]
forall a. [a] -> [a] -> [a]
++ TyConMap -> [(Type, Type)] -> [(TyVar, Type)]
solveNonAbsurds TyConMap
tcm [(Type, Type)]
eqs
 where
  solvers :: [(Type, Type) -> [TypeEqSolution]]
solvers = [TypeEqSolution -> [TypeEqSolution]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TypeEqSolution -> [TypeEqSolution])
-> ((Type, Type) -> TypeEqSolution)
-> (Type, Type)
-> [TypeEqSolution]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type, Type) -> TypeEqSolution
solveAdd, TyConMap -> (Type, Type) -> [TypeEqSolution]
solveEq TyConMap
tcm]
  solved :: [(TyVar, Type)]
solved = [TypeEqSolution] -> [(TyVar, Type)]
catSolutions ([[TypeEqSolution]] -> [TypeEqSolution]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [(Type, Type) -> [TypeEqSolution]
s (Type, Type)
eq | (Type, Type) -> [TypeEqSolution]
s <- [(Type, Type) -> [TypeEqSolution]]
solvers])

-- | Solve simple equalities such as:
--
--   * a ~ 3
--   * 3 ~ a
--   * SomeType a b ~ SomeType 3 5
--   * SomeType 3 5 ~ SomeType a b
--   * SomeType a 5 ~ SomeType 3 b
--
solveEq :: TyConMap -> (Type, Type) -> [TypeEqSolution]
solveEq :: TyConMap -> (Type, Type) -> [TypeEqSolution]
solveEq tcm :: TyConMap
tcm (TyConMap -> Type -> Type
coreView TyConMap
tcm -> Type
left, TyConMap -> Type -> Type
coreView TyConMap
tcm -> Type
right) =
  case (Type
left, Type
right) of
    (VarTy tyVar :: TyVar
tyVar, ConstTy {}) ->
      -- a ~ 3
      [(TyVar, Type) -> TypeEqSolution
Solution (TyVar
tyVar, Type
right)]
    (ConstTy {}, VarTy tyVar :: TyVar
tyVar) ->
      -- 3 ~ a
      [(TyVar, Type) -> TypeEqSolution
Solution (TyVar
tyVar, Type
left)]
    (ConstTy {}, ConstTy {}) ->
      -- Int /= Char
      if Type
left Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
/= Type
right then [TypeEqSolution
AbsurdSolution] else []
    (LitTy {}, LitTy {}) ->
      -- 3 /= 5
      if Type
left Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
/= Type
right then [TypeEqSolution
AbsurdSolution] else []
    _ ->
      -- The call to 'coreView' at the start of 'solveEq' should have reduced
      -- all solvable type families. If we encounter one here that means the
      -- type family is stuck (and that we shouldn't compare it to anything!).
      if (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (TyConMap -> Type -> Bool
isTypeFamilyApplication TyConMap
tcm) [Type
left, Type
right] then
        []
      else
        case (Type -> TypeView
tyView Type
left, Type -> TypeView
tyView Type
right) of
          (TyConApp leftNm :: TyConName
leftNm leftTys :: [Type]
leftTys, TyConApp rightNm :: TyConName
rightNm rightTys :: [Type]
rightTys) ->
            -- SomeType a b ~ SomeType 3 5 (or other way around)
            if TyConName
leftNm TyConName -> TyConName -> Bool
forall a. Eq a => a -> a -> Bool
== TyConName
rightNm then
              [[TypeEqSolution]] -> [TypeEqSolution]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (((Type, Type) -> [TypeEqSolution])
-> [(Type, Type)] -> [[TypeEqSolution]]
forall a b. (a -> b) -> [a] -> [b]
map (TyConMap -> (Type, Type) -> [TypeEqSolution]
solveEq TyConMap
tcm) ([Type] -> [Type] -> [(Type, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
leftTys [Type]
rightTys))
            else
              [TypeEqSolution
AbsurdSolution]
          _ ->
            []

-- | Solve equations supported by @normalizeAdd@. See documentation of
-- @TypeEqSolution@ to understand the return value.
solveAdd
  :: (Type, Type)
  -> TypeEqSolution
solveAdd :: (Type, Type) -> TypeEqSolution
solveAdd ab :: (Type, Type)
ab =
  case (Type, Type) -> Maybe (Integer, Integer, Type)
normalizeAdd (Type, Type)
ab of
    Just (n :: Integer
n, m :: Integer
m, VarTy tyVar :: TyVar
tyVar) ->
      if Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= 0 Bool -> Bool -> Bool
&& Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= 0 Bool -> Bool -> Bool
&& Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= 0 then
        (TyVar, Type) -> TypeEqSolution
Solution (TyVar
tyVar, (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
m))))
      else
        TypeEqSolution
AbsurdSolution
    _ ->
      TypeEqSolution
NoSolution

-- | If type is an equation, return LHS and RHS.
typeEq
  :: TyConMap
  -> Type
  -> Maybe (Type, Type)
typeEq :: TyConMap -> Type -> Maybe (Type, Type)
typeEq tcm :: TyConMap
tcm ty :: Type
ty =
 case Type -> TypeView
tyView (TyConMap -> Type -> Type
coreView TyConMap
tcm Type
ty) of
  TyConApp (TyConName -> OccName
forall a. Name a -> OccName
nameOcc -> OccName
"GHC.Prim.~#") [_, _, left :: Type
left, right :: Type
right] ->
    (Type, Type) -> Maybe (Type, Type)
forall a. a -> Maybe a
Just (TyConMap -> Type -> Type
coreView TyConMap
tcm Type
left, TyConMap -> Type -> Type
coreView TyConMap
tcm Type
right)
  _ ->
    Maybe (Type, Type)
forall a. Maybe a
Nothing

-- | Get constraint equations
altEqs
  :: TyConMap
  -> Alt
  -> [(Type, Type)]
altEqs :: TyConMap -> Alt -> [(Type, Type)]
altEqs tcm :: TyConMap
tcm (pat :: Pat
pat, _term :: Term
_term) =
 [Maybe (Type, Type)] -> [(Type, Type)]
forall a. [Maybe a] -> [a]
catMaybes ((Var Term -> Maybe (Type, Type))
-> [Var Term] -> [Maybe (Type, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (TyConMap -> Type -> Maybe (Type, Type)
typeEq TyConMap
tcm (Type -> Maybe (Type, Type))
-> (Var Term -> Type) -> Var Term -> Maybe (Type, Type)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Var Term -> Type
forall a. Var a -> Type
varType) (([TyVar], [Var Term]) -> [Var Term]
forall a b. (a, b) -> b
snd (Pat -> ([TyVar], [Var Term])
patIds Pat
pat)))

-- | Tests for unreachable alternative due to types being "absurd". See
-- @isAbsurdEq@ for more info.
isAbsurdAlt
  :: TyConMap
  -> Alt
  -> Bool
isAbsurdAlt :: TyConMap -> Alt -> Bool
isAbsurdAlt tcm :: TyConMap
tcm alt :: Alt
alt =
  ((Type, Type) -> Bool) -> [(Type, Type)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (TyConMap -> (Type, Type) -> Bool
isAbsurdEq TyConMap
tcm) (TyConMap -> Alt -> [(Type, Type)]
altEqs TyConMap
tcm Alt
alt)

-- | Determines if an "equation" obtained through @altEqs@ or @typeEq@ is
-- absurd. That is, it tests if two types that are definitely not equal are
-- asserted to be equal OR if the computation of the types yield some absurd
-- (intermediate) result such as -1.
isAbsurdEq
  :: TyConMap
  -> (Type, Type)
  -> Bool
isAbsurdEq :: TyConMap -> (Type, Type) -> Bool
isAbsurdEq tcm :: TyConMap
tcm ((left0 :: Type
left0, right0 :: Type
right0)) =
  case (TyConMap -> Type -> Type
coreView TyConMap
tcm Type
left0, TyConMap -> Type -> Type
coreView TyConMap
tcm Type
right0) of
    ((Type, Type) -> TypeEqSolution
solveAdd -> TypeEqSolution
AbsurdSolution) -> Bool
True
    lr :: (Type, Type)
lr -> (TypeEqSolution -> Bool) -> [TypeEqSolution] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (TypeEqSolution -> TypeEqSolution -> Bool
forall a. Eq a => a -> a -> Bool
==TypeEqSolution
AbsurdSolution) (TyConMap -> (Type, Type) -> [TypeEqSolution]
solveEq TyConMap
tcm (Type, Type)
lr)

-- Safely substitute global type variables in a list of potentially
-- shadowing type variables.
substGlobalsInExistentials
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> [(TyVar, Type)]
  -- ^ Substitutions
  -> [TyVar]
substGlobalsInExistentials :: InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials is :: InScopeSet
is exts :: [TyVar]
exts substs0 :: [(TyVar, Type)]
substs0 = [TyVar]
result
  -- TODO: Is is actually possible that existentials shadow each other? If they
  -- TODO: can't, we can remove this function
  where
    iss :: [InScopeSet]
iss     = (InScopeSet -> TyVar -> InScopeSet)
-> InScopeSet -> [TyVar] -> [InScopeSet]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl InScopeSet -> TyVar -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
is [TyVar]
exts
    substs1 :: [Subst]
substs1 = (InScopeSet -> Subst) -> [InScopeSet] -> [Subst]
forall a b. (a -> b) -> [a] -> [b]
map (\is_ :: InScopeSet
is_ -> Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (InScopeSet -> Subst
mkSubst InScopeSet
is_) [(TyVar, Type)]
substs0) [InScopeSet]
iss
    result :: [TyVar]
result  = (Subst -> TyVar -> TyVar) -> [Subst] -> [TyVar] -> [TyVar]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Subst -> TyVar -> TyVar
forall a. HasCallStack => Subst -> Var a -> Var a
substTyInVar [Subst]
substs1 [TyVar]
exts

-- | Safely substitute type variables in a list of existentials. This function
-- will account for cases where existentials shadow each other.
substInExistentialsList
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> [(TyVar, Type)]
  -- ^ Substitutions
  -> [TyVar]
substInExistentialsList :: InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substInExistentialsList is :: InScopeSet
is exts :: [TyVar]
exts substs :: [(TyVar, Type)]
substs =
  ([TyVar] -> (TyVar, Type) -> [TyVar])
-> [TyVar] -> [(TyVar, Type)] -> [TyVar]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (HasCallStack => InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
substInExistentials InScopeSet
is) [TyVar]
exts [(TyVar, Type)]
substs

-- | Safely substitute a type variable in a list of existentials. This function
-- will account for cases where existentials shadow each other.
substInExistentials
  :: HasCallStack
  => InScopeSet
  -- ^ Variables in scope
  -> [TyVar]
  -- ^ List of existentials to apply the substitution for
  -> (TyVar, Type)
  -- ^ Substitution
  -> [TyVar]
substInExistentials :: InScopeSet -> [TyVar] -> (TyVar, Type) -> [TyVar]
substInExistentials is :: InScopeSet
is exts :: [TyVar]
exts subst :: (TyVar, Type)
subst@(typeVar :: TyVar
typeVar, _type :: Type
_type) =
  -- TODO: Is is actually possible that existentials shadow each other? If they
  -- TODO: can't, we can remove this function
  case TyVar -> [TyVar] -> [Int]
forall a. Eq a => a -> [a] -> [Int]
elemIndices TyVar
typeVar [TyVar]
exts of
    [] ->
      -- We're not replacing any of the existentials, but a global variable
      HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is [TyVar]
exts [(TyVar, Type)
subst]
    ([Int] -> Int
forall a. [a] -> a
last -> Int
i) ->
      -- We're replacing an existential. That means we're not touching any
      -- variables that were introduced before it. For all variables after it,
      -- it is as we would replace global variables in them.
      Int -> [TyVar] -> [TyVar]
forall a. Int -> [a] -> [a]
take (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) [TyVar]
exts [TyVar] -> [TyVar] -> [TyVar]
forall a. [a] -> [a] -> [a]
++ HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is (Int -> [TyVar] -> [TyVar]
forall a. Int -> [a] -> [a]
drop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) [TyVar]
exts) [(TyVar, Type)
subst]

-- | Determine the type of a term
termType
  :: TyConMap
  -> Term
  -> Type
termType :: TyConMap -> Term -> Type
termType m :: TyConMap
m e :: Term
e = case Term
e of
  Var t :: Var Term
t          -> Var Term -> Type
forall a. Var a -> Type
varType Var Term
t
  Data dc :: DataCon
dc        -> DataCon -> Type
dcType DataCon
dc
  Literal l :: Literal
l      -> Literal -> Type
literalType Literal
l
  Prim _ t :: PrimInfo
t       -> PrimInfo -> Type
primType PrimInfo
t
  Lam v :: Var Term
v e' :: Term
e'       -> Type -> Type -> Type
mkFunTy (Var Term -> Type
forall a. Var a -> Type
varType Var Term
v) (TyConMap -> Term -> Type
termType TyConMap
m Term
e')
  TyLam tv :: TyVar
tv e' :: Term
e'    -> TyVar -> Type -> Type
ForAllTy TyVar
tv (TyConMap -> Term -> Type
termType TyConMap
m Term
e')
  App _ _        -> case Term -> (Term, [Either Term Type])
collectArgs Term
e of
                      (fun :: Term
fun, args :: [Either Term Type]
args) -> Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs Term
e TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
fun) [Either Term Type]
args
  TyApp _ _      -> case Term -> (Term, [Either Term Type])
collectArgs Term
e of
                      (fun :: Term
fun, args :: [Either Term Type]
args) -> Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs Term
e TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
fun) [Either Term Type]
args
  Letrec _ e' :: Term
e'    -> TyConMap -> Term -> Type
termType TyConMap
m Term
e'
  Case _ ty :: Type
ty _    -> Type
ty
  Cast _ _ ty2 :: Type
ty2   -> Type
ty2
  Tick _ e' :: Term
e'      -> TyConMap -> Term -> Type
termType TyConMap
m Term
e'

-- | Split a (Type)Abstraction in the bound variables and the abstracted term
collectBndrs :: Term
             -> ([Either Id TyVar], Term)
collectBndrs :: Term -> ([Either (Var Term) TyVar], Term)
collectBndrs = [Either (Var Term) TyVar]
-> Term -> ([Either (Var Term) TyVar], Term)
go []
  where
    go :: [Either (Var Term) TyVar]
-> Term -> ([Either (Var Term) TyVar], Term)
go bs :: [Either (Var Term) TyVar]
bs (Lam v :: Var Term
v e' :: Term
e')    = [Either (Var Term) TyVar]
-> Term -> ([Either (Var Term) TyVar], Term)
go (Var Term -> Either (Var Term) TyVar
forall a b. a -> Either a b
Left Var Term
vEither (Var Term) TyVar
-> [Either (Var Term) TyVar] -> [Either (Var Term) TyVar]
forall a. a -> [a] -> [a]
:[Either (Var Term) TyVar]
bs) Term
e'
    go bs :: [Either (Var Term) TyVar]
bs (TyLam tv :: TyVar
tv e' :: Term
e') = [Either (Var Term) TyVar]
-> Term -> ([Either (Var Term) TyVar], Term)
go (TyVar -> Either (Var Term) TyVar
forall a b. b -> Either a b
Right TyVar
tvEither (Var Term) TyVar
-> [Either (Var Term) TyVar] -> [Either (Var Term) TyVar]
forall a. a -> [a] -> [a]
:[Either (Var Term) TyVar]
bs) Term
e'
    go bs :: [Either (Var Term) TyVar]
bs e' :: Term
e'            = ([Either (Var Term) TyVar] -> [Either (Var Term) TyVar]
forall a. [a] -> [a]
reverse [Either (Var Term) TyVar]
bs,Term
e')

-- | Get the result type of a polymorphic function given a list of arguments
applyTypeToArgs
  :: Term
  -> TyConMap
  -> Type
  -> [Either Term Type]
  -> Type
applyTypeToArgs :: Term -> TyConMap -> Type -> [Either Term Type] -> Type
applyTypeToArgs e :: Term
e m :: TyConMap
m opTy :: Type
opTy args :: [Either Term Type]
args = Type -> [Either Term Type] -> Type
forall a. Type -> [Either a Type] -> Type
go Type
opTy [Either Term Type]
args
 where
  go :: Type -> [Either a Type] -> Type
go opTy' :: Type
opTy' []               = Type
opTy'
  go opTy' :: Type
opTy' (Right ty :: Type
ty:args' :: [Either a Type]
args') = Type -> [Type] -> [Either a Type] -> Type
goTyArgs Type
opTy' [Type
ty] [Either a Type]
args'
  go opTy' :: Type
opTy' (Left _:args' :: [Either a Type]
args')   = case TyConMap -> Type -> Maybe (Type, Type)
splitFunTy TyConMap
m Type
opTy' of
    Just (_,resTy :: Type
resTy) -> Type -> [Either a Type] -> Type
go Type
resTy [Either a Type]
args'
    _ -> String -> Type
forall a. HasCallStack => String -> a
error (String -> Type) -> String -> Type
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines ["applyTypeToArgs:"
                         ,"Expression: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Term -> String
forall p. PrettyPrec p => p -> String
showPpr Term
e
                         ,"Type: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall p. PrettyPrec p => p -> String
showPpr Type
opTy
                         ,"Args: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines ((Either Term Type -> String) -> [Either Term Type] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((Term -> String) -> (Type -> String) -> Either Term Type -> String
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Term -> String
forall p. PrettyPrec p => p -> String
showPpr Type -> String
forall p. PrettyPrec p => p -> String
showPpr) [Either Term Type]
args)
                         ]

  goTyArgs :: Type -> [Type] -> [Either a Type] -> Type
goTyArgs opTy' :: Type
opTy' revTys :: [Type]
revTys (Right ty :: Type
ty:args' :: [Either a Type]
args') = Type -> [Type] -> [Either a Type] -> Type
goTyArgs Type
opTy' (Type
tyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
revTys) [Either a Type]
args'
  goTyArgs opTy' :: Type
opTy' revTys :: [Type]
revTys args' :: [Either a Type]
args'            = Type -> [Either a Type] -> Type
go (TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
opTy' ([Type] -> [Type]
forall a. [a] -> [a]
reverse [Type]
revTys)) [Either a Type]
args'

-- | Like 'piResultTyMaybe', but errors out when a type application is not
-- valid.
--
-- Do not iterate 'piResultTy', because it's inefficient to substitute one
-- variable at a time; instead use 'piResultTys'
piResultTy
  :: TyConMap
  -> Type
  -> Type
  -> Type
piResultTy :: TyConMap -> Type -> Type -> Type
piResultTy m :: TyConMap
m ty :: Type
ty arg :: Type
arg = case TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe TyConMap
m Type
ty Type
arg of
  Just res :: Type
res -> Type
res
  Nothing  -> String -> Doc ClashAnnotation -> Type
forall ann a. String -> Doc ann -> a
pprPanic "piResultTy" (Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
ty Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Doc ClashAnnotation
forall ann. Doc ann
line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
arg)

-- | Like 'piResultTys' but for a single argument.
--
-- Do not iterate 'piResultTyMaybe', because it's inefficient to substitute one
-- variable at a time; instead use 'piResultTys'
piResultTyMaybe
  :: TyConMap
  -> Type
  -> Type
  -> Maybe Type
piResultTyMaybe :: TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe m :: TyConMap
m ty :: Type
ty arg :: Type
arg
  | Just ty' :: Type
ty' <- TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m Type
ty
  = TyConMap -> Type -> Type -> Maybe Type
piResultTyMaybe TyConMap
m Type
ty' Type
arg
  | FunTy _ res :: Type
res <- Type -> TypeView
tyView Type
ty
  = Type -> Maybe Type
forall a. a -> Maybe a
Just Type
res
  | ForAllTy tv :: TyVar
tv res :: Type
res <- Type
ty
  = let emptySubst :: Subst
emptySubst = InScopeSet -> Subst
mkSubst (VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type
arg,Type
res]))
    in  Type -> Maybe Type
forall a. a -> Maybe a
Just (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy (Subst -> TyVar -> Type -> Subst
extendTvSubst Subst
emptySubst TyVar
tv Type
arg) Type
res)
  | Bool
otherwise
  = Maybe Type
forall a. Maybe a
Nothing

-- | @(piResultTys f_ty [ty1, ..., tyn])@ give sthe type of @(f ty1 .. tyn)@
-- where @f :: f_ty@
--
-- 'piResultTys' is interesting because:
--
--    1. 'f_ty' may have more foralls than there are args
--    2. Less obviously, it may have fewer foralls
--
-- Fore case 2. think of:
--
--   piResultTys (forall a . a) [forall b.b, Int]
--
-- This really can happen, such as situations involving 'undefined's type:
--
--   undefined :: forall a. a
--
--   undefined (forall b. b -> b) Int
--
-- This term should have the type @(Int -> Int)@, but notice that there are
-- more type args than foralls in 'undefined's type.
--
-- For efficiency reasons, when there are no foralls, we simply drop arrows from
-- a function type/kind.
piResultTys
  :: TyConMap
  -> Type
  -> [Type]
  -> Type
piResultTys :: TyConMap -> Type -> [Type] -> Type
piResultTys _ ty :: Type
ty [] = Type
ty
piResultTys m :: TyConMap
m ty :: Type
ty origArgs :: [Type]
origArgs@(arg :: Type
arg:args :: [Type]
args)
  | Just ty' :: Type
ty' <- TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m Type
ty
  = TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
ty' [Type]
origArgs
  | FunTy _ res :: Type
res <- Type -> TypeView
tyView Type
ty
  = TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
res [Type]
args
  | ForAllTy tv :: TyVar
tv res :: Type
res <- Type
ty
  = VarEnv Type -> Type -> [Type] -> Type
go (TyVar -> Type -> VarEnv Type -> VarEnv Type
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv Type
arg VarEnv Type
forall a. VarEnv a
emptyVarEnv) Type
res [Type]
args
  | Bool
otherwise
  = String -> Doc ClashAnnotation -> Type
forall ann a. String -> Doc ann -> a
pprPanic "piResultTys1" (Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
ty Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Doc ClashAnnotation
forall ann. Doc ann
line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> [Type] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [Type]
origArgs)
 where
  inScope :: InScopeSet
inScope = VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes (Type
tyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
origArgs))

  go :: VarEnv Type -> Type -> [Type] -> Type
go env :: VarEnv Type
env ty' :: Type
ty' [] = HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy (InScopeSet -> VarEnv Type -> Subst
mkTvSubst InScopeSet
inScope VarEnv Type
env) Type
ty'
  go env :: VarEnv Type
env ty' :: Type
ty' allArgs :: [Type]
allArgs@(arg' :: Type
arg':args' :: [Type]
args')
    | Just ty'' :: Type
ty'' <- TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m Type
ty'
    = VarEnv Type -> Type -> [Type] -> Type
go VarEnv Type
env Type
ty'' [Type]
allArgs
    | FunTy _ res :: Type
res <- Type -> TypeView
tyView Type
ty'
    = VarEnv Type -> Type -> [Type] -> Type
go VarEnv Type
env Type
res [Type]
args'
    | ForAllTy tv :: TyVar
tv res :: Type
res <- Type
ty'
    = VarEnv Type -> Type -> [Type] -> Type
go (TyVar -> Type -> VarEnv Type -> VarEnv Type
forall b a. Var b -> a -> VarEnv a -> VarEnv a
extendVarEnv TyVar
tv Type
arg' VarEnv Type
env) Type
res [Type]
args'
    | VarTy tv :: TyVar
tv <- Type
ty'
    , Just ty'' :: Type
ty'' <- TyVar -> VarEnv Type -> Maybe Type
forall b a. Var b -> VarEnv a -> Maybe a
lookupVarEnv TyVar
tv VarEnv Type
env
      -- Deals with (piResultTys  (forall a.a) [forall b.b, Int])
    = TyConMap -> Type -> [Type] -> Type
piResultTys TyConMap
m Type
ty'' [Type]
allArgs
    | Bool
otherwise
    = String -> Doc ClashAnnotation -> Type
forall ann a. String -> Doc ann -> a
pprPanic "piResultTys2" (Type -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr Type
ty' Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Doc ClashAnnotation
forall ann. Doc ann
line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> [Type] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [Type]
origArgs Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> Doc ClashAnnotation
forall ann. Doc ann
line Doc ClashAnnotation -> Doc ClashAnnotation -> Doc ClashAnnotation
forall a. Semigroup a => a -> a -> a
<> [Type] -> Doc ClashAnnotation
forall p. PrettyPrec p => p -> Doc ClashAnnotation
ppr [Type]
allArgs)

-- | Get the list of term-binders out of a DataType pattern
patIds :: Pat -> ([TyVar],[Id])
patIds :: Pat -> ([TyVar], [Var Term])
patIds (DataPat _  tvs :: [TyVar]
tvs ids :: [Var Term]
ids) = ([TyVar]
tvs,[Var Term]
ids)
patIds _                    = ([],[])

patVars :: Pat -> [Var a]
patVars :: Pat -> [Var a]
patVars (DataPat _ tvs :: [TyVar]
tvs ids :: [Var Term]
ids) = [TyVar] -> [Var a]
forall a b. Coercible a b => a -> b
coerce [TyVar]
tvs [Var a] -> [Var a] -> [Var a]
forall a. [a] -> [a] -> [a]
++ [Var Term] -> [Var a]
forall a b. Coercible a b => a -> b
coerce [Var Term]
ids
patVars _ = []

-- | Abstract a term over a list of term and type variables
mkAbstraction :: Term
              -> [Either Id TyVar]
              -> Term
mkAbstraction :: Term -> [Either (Var Term) TyVar] -> Term
mkAbstraction = (Either (Var Term) TyVar -> Term -> Term)
-> Term -> [Either (Var Term) TyVar] -> Term
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Var Term -> Term -> Term)
-> (TyVar -> Term -> Term)
-> Either (Var Term) TyVar
-> Term
-> Term
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Var Term -> Term -> Term
Lam TyVar -> Term -> Term
TyLam)

-- | Abstract a term over a list of term variables
mkTyLams :: Term
         -> [TyVar]
         -> Term
mkTyLams :: Term -> [TyVar] -> Term
mkTyLams tm :: Term
tm = Term -> [Either (Var Term) TyVar] -> Term
mkAbstraction Term
tm ([Either (Var Term) TyVar] -> Term)
-> ([TyVar] -> [Either (Var Term) TyVar]) -> [TyVar] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TyVar -> Either (Var Term) TyVar)
-> [TyVar] -> [Either (Var Term) TyVar]
forall a b. (a -> b) -> [a] -> [b]
map TyVar -> Either (Var Term) TyVar
forall a b. b -> Either a b
Right

-- | Abstract a term over a list of type variables
mkLams :: Term
       -> [Id]
       -> Term
mkLams :: Term -> [Var Term] -> Term
mkLams tm :: Term
tm = Term -> [Either (Var Term) TyVar] -> Term
mkAbstraction Term
tm ([Either (Var Term) TyVar] -> Term)
-> ([Var Term] -> [Either (Var Term) TyVar]) -> [Var Term] -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Var Term -> Either (Var Term) TyVar)
-> [Var Term] -> [Either (Var Term) TyVar]
forall a b. (a -> b) -> [a] -> [b]
map Var Term -> Either (Var Term) TyVar
forall a b. a -> Either a b
Left

-- | Apply a list of types and terms to a term
mkApps :: Term
       -> [Either Term Type]
       -> Term
mkApps :: Term -> [Either Term Type] -> Term
mkApps = (Term -> Either Term Type -> Term)
-> Term -> [Either Term Type] -> Term
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\e :: Term
e a :: Either Term Type
a -> (Term -> Term) -> (Type -> Term) -> Either Term Type -> Term
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Term -> Term -> Term
App Term
e) (Term -> Type -> Term
TyApp Term
e) Either Term Type
a)

-- | Apply a list of terms to a term
mkTmApps :: Term
         -> [Term]
         -> Term
mkTmApps :: Term -> [Term] -> Term
mkTmApps = (Term -> Term -> Term) -> Term -> [Term] -> Term
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Term -> Term -> Term
App

-- | Apply a list of types to a term
mkTyApps :: Term
         -> [Type]
         -> Term
mkTyApps :: Term -> [Type] -> Term
mkTyApps = (Term -> Type -> Term) -> Term -> [Type] -> Term
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Term -> Type -> Term
TyApp

mkTicks
  :: Term
  -> [TickInfo]
  -> Term
mkTicks :: Term -> [TickInfo] -> Term
mkTicks tm :: Term
tm ticks :: [TickInfo]
ticks = (Term -> TickInfo -> Term) -> Term -> [TickInfo] -> Term
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\e :: Term
e s :: TickInfo
s -> TickInfo -> Term -> Term
Tick TickInfo
s Term
e) Term
tm ([TickInfo] -> [TickInfo]
forall a. Eq a => [a] -> [a]
nub [TickInfo]
ticks)

-- | Does a term have a function type?
isFun :: TyConMap
      -> Term
      -> Bool
isFun :: TyConMap -> Term -> Bool
isFun m :: TyConMap
m t :: Term
t = TyConMap -> Type -> Bool
isFunTy TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
t)

-- | Does a term have a function or polymorphic type?
isPolyFun :: TyConMap
          -> Term
          -> Bool
isPolyFun :: TyConMap -> Term -> Bool
isPolyFun m :: TyConMap
m t :: Term
t = TyConMap -> Type -> Bool
isPolyFunCoreTy TyConMap
m (TyConMap -> Term -> Type
termType TyConMap
m Term
t)

-- | Is a term a term-abstraction?
isLam :: Term
      -> Bool
isLam :: Term -> Bool
isLam (Lam {}) = Bool
True
isLam _        = Bool
False

-- | Is a term a recursive let-binding?
isLet :: Term
      -> Bool
isLet :: Term -> Bool
isLet (Letrec {}) = Bool
True
isLet _           = Bool
False

-- | Is a term a variable reference?
isVar :: Term
      -> Bool
isVar :: Term -> Bool
isVar (Var {}) = Bool
True
isVar _        = Bool
False

isLocalVar
  :: Term
  -> Bool
isLocalVar :: Term -> Bool
isLocalVar (Var v :: Var Term
v) = Var Term -> Bool
forall a. Var a -> Bool
isLocalId Var Term
v
isLocalVar _ = Bool
False

-- | Is a term a datatype constructor?
isCon :: Term
      -> Bool
isCon :: Term -> Bool
isCon (Data {}) = Bool
True
isCon _         = Bool
False

-- | Is a term a primitive?
isPrim :: Term
       -> Bool
isPrim :: Term -> Bool
isPrim (Prim {}) = Bool
True
isPrim _         = Bool
False

-- | Make variable reference out of term variable
idToVar :: Id
        -> Term
idToVar :: Var Term -> Term
idToVar i :: Var Term
i@(Id {}) = Var Term -> Term
Var Var Term
i
idToVar tv :: Var Term
tv        = String -> Term
forall a. HasCallStack => String -> a
error (String -> Term) -> String -> Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "idToVar: tyVar: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Var Term -> String
forall p. PrettyPrec p => p -> String
showPpr Var Term
tv

-- | Make a term variable out of a variable reference
varToId :: Term
        -> Id
varToId :: Term -> Var Term
varToId (Var i :: Var Term
i) = Var Term
i
varToId e :: Term
e       = String -> Var Term
forall a. HasCallStack => String -> a
error (String -> Var Term) -> String -> Var Term
forall a b. (a -> b) -> a -> b
$ $(curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "varToId: not a var: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Term -> String
forall p. PrettyPrec p => p -> String
showPpr Term
e

termSize :: Term
         -> Word
termSize :: Term -> Word
termSize (Var {})     = 1
termSize (Data {})    = 1
termSize (Literal {}) = 1
termSize (Prim {})    = 1
termSize (Lam _ e :: Term
e)    = Term -> Word
termSize Term
e Word -> Word -> Word
forall a. Num a => a -> a -> a
+ 1
termSize (TyLam _ e :: Term
e)  = Term -> Word
termSize Term
e
termSize (App e1 :: Term
e1 e2 :: Term
e2)  = Term -> Word
termSize Term
e1 Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Term -> Word
termSize Term
e2
termSize (TyApp e :: Term
e _)  = Term -> Word
termSize Term
e
termSize (Cast e :: Term
e _ _) = Term -> Word
termSize Term
e
termSize (Tick _ e :: Term
e)   = Term -> Word
termSize Term
e
termSize (Letrec bndrs :: [LetBinding]
bndrs e :: Term
e) = [Word] -> Word
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Word
bodySzWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
bndrSzs)
 where
  bndrSzs :: [Word]
bndrSzs = (LetBinding -> Word) -> [LetBinding] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Word
termSize (Term -> Word) -> (LetBinding -> Term) -> LetBinding -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LetBinding -> Term
forall a b. (a, b) -> b
snd) [LetBinding]
bndrs
  bodySz :: Word
bodySz  = Term -> Word
termSize Term
e
termSize (Case subj :: Term
subj _ alts :: [Alt]
alts) = [Word] -> Word
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Word
subjSzWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
altSzs)
 where
  subjSz :: Word
subjSz = Term -> Word
termSize Term
subj
  altSzs :: [Word]
altSzs = (Alt -> Word) -> [Alt] -> [Word]
forall a b. (a -> b) -> [a] -> [b]
map (Term -> Word
termSize (Term -> Word) -> (Alt -> Term) -> Alt -> Word
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Alt -> Term
forall a b. (a, b) -> b
snd) [Alt]
alts

-- | Create a vector of supplied elements
mkVec :: DataCon -- ^ The Nil constructor
      -> DataCon -- ^ The Cons (:>) constructor
      -> Type    -- ^ Element type
      -> Integer -- ^ Length of the vector
      -> [Term]  -- ^ Elements to put in the vector
      -> Term
mkVec :: DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkVec nilCon :: DataCon
nilCon consCon :: DataCon
consCon resTy :: Type
resTy = Integer -> [Term] -> Term
go
  where
    go :: Integer -> [Term] -> Term
go _ [] = Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
nilCon) [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy 0))
                                   ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                                   ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Type -> Term
primCo Type
nilCoTy)
                                   ]

    go n :: Integer
n (x :: Term
x:xs :: [Term]
xs) = Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
consCon) [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                        ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                                        ,Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Type -> Term
primCo (Integer -> Type
consCoTy Integer
n))
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
x
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> [Term] -> Term
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Term]
xs)]

    nilCoTy :: Type
nilCoTy    = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Type] -> [Type]) -> Maybe [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$! DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
nilCon  [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy 0))
                                                             ,Type
resTy])
    consCoTy :: Integer -> Type
consCoTy n :: Integer
n = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Type] -> [Type]) -> Maybe [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$! DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon
                                                     [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                                     ,Type
resTy
                                                     ,(LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))])

-- | Append elements to the supplied vector
appendToVec :: DataCon -- ^ The Cons (:>) constructor
            -> Type    -- ^ Element type
            -> Term    -- ^ The vector to append the elements to
            -> Integer -- ^ Length of the vector
            -> [Term]  -- ^ Elements to append
            -> Term
appendToVec :: DataCon -> Type -> Term -> Integer -> [Term] -> Term
appendToVec consCon :: DataCon
consCon resTy :: Type
resTy vec :: Term
vec = Integer -> [Term] -> Term
go
  where
    go :: Integer -> [Term] -> Term
go _ []     = Term
vec
    go n :: Integer
n (x :: Term
x:xs :: [Term]
xs) = Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
consCon) [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                        ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                                        ,Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Type -> Term
primCo (Integer -> Type
consCoTy Integer
n))
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left Term
x
                                        ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> [Term] -> Term
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Term]
xs)]

    consCoTy :: Integer -> Type
consCoTy n :: Integer
n = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Type] -> [Type]) -> Maybe [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$! DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon
                                                   [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                                   ,Type
resTy
                                                   ,(LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))])


availableUniques
  :: Term
  -> [Unique]
availableUniques :: Term -> [Int]
availableUniques t :: Term
t = [ Int
n | Int
n <- [0..] , Int
n Int -> IntSet -> Bool
`IntSet.notMember` IntSet
avoid ]
 where
  avoid :: IntSet
avoid = Getting (IntSet -> IntSet) Term (Var Any)
-> (Var Any -> IntSet -> IntSet) -> Term -> IntSet -> IntSet
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting (IntSet -> IntSet) Term (Var Any)
forall a. Fold Term (Var a)
termFreeVarsX (\a :: Var Any
a i :: IntSet
i -> Int -> IntSet -> IntSet
IntSet.insert (Var Any -> Int
forall a. Var a -> Int
varUniq Var Any
a) IntSet
i) Term
t
            IntSet
IntSet.empty

-- | Create let-bindings with case-statements that select elements out of a
-- vector. Returns both the variables to which element-selections are bound
-- and the let-bindings
extractElems
  :: Supply
  -- ^ Unique supply
  -> InScopeSet
  -- ^ (Superset of) in scope variables
  -> DataCon
  -- ^ The Cons (:>) constructor
  -> Type
  -- ^ The element type
  -> Char
  -- ^ Char to append to the bound variable names
  -> Integer
  -- ^ Length of the vector
  -> Term
  -- ^ The vector
  -> (Supply, [(Term,[LetBinding])])
extractElems :: Supply
-> InScopeSet
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, [(Term, [LetBinding])])
extractElems supply :: Supply
supply inScope :: InScopeSet
inScope consCon :: DataCon
consCon resTy :: Type
resTy s :: Char
s maxN :: Integer
maxN vec :: Term
vec =
  ((Supply, InScopeSet) -> Supply)
-> ((Supply, InScopeSet), [(Term, [LetBinding])])
-> (Supply, [(Term, [LetBinding])])
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Supply, InScopeSet) -> Supply
forall a b. (a, b) -> a
fst (Integer
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), [(Term, [LetBinding])])
go Integer
maxN (Supply
supply,InScopeSet
inScope) Term
vec)
 where
  go :: Integer -> (Supply,InScopeSet) -> Term
     -> ((Supply,InScopeSet),[(Term,[LetBinding])])
  go :: Integer
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), [(Term, [LetBinding])])
go 0 uniqs :: (Supply, InScopeSet)
uniqs _ = ((Supply, InScopeSet)
uniqs,[])
  go n :: Integer
n uniqs0 :: (Supply, InScopeSet)
uniqs0 e :: Term
e =
    ((Supply, InScopeSet)
uniqs3,(Term
elNVar,[(Var Term
elNId, Term
lhs),(Var Term
restNId, Term
rhs)])(Term, [LetBinding])
-> [(Term, [LetBinding])] -> [(Term, [LetBinding])]
forall a. a -> [a] -> [a]
:[(Term, [LetBinding])]
restVs)
   where
    tys :: [Type]
tys = [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n)),Type
resTy,(LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))]
    (Just idTys :: [Type]
idTys) = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
consCon [Type]
tys
    restTy :: Type
restTy       = [Type] -> Type
forall a. [a] -> a
last [Type]
idTys

    (uniqs1 :: (Supply, InScopeSet)
uniqs1,mTV :: TyVar
mTV) = (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar (Supply, InScopeSet)
uniqs0 ("m",Type
typeNatKind)
    (uniqs2 :: (Supply, InScopeSet)
uniqs2,[elNId :: Var Term
elNId,restNId :: Var Term
restNId,co :: Var Term
co,el :: Var Term
el,rest :: Var Term
rest]) =
      ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var Term))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Var Term])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var Term)
mkUniqSystemId (Supply, InScopeSet)
uniqs1 ([(OccName, Type)] -> ((Supply, InScopeSet), [Var Term]))
-> [(OccName, Type)] -> ((Supply, InScopeSet), [Var Term])
forall a b. (a -> b) -> a -> b
$ [OccName] -> [Type] -> [(OccName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip
        ["el" OccName -> OccName -> OccName
`T.append` (Char
s Char -> OccName -> OccName
`T.cons` String -> OccName
T.pack (Integer -> String
forall a. Show a => a -> String
show (Integer
maxNInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
n)))
        ,"rest" OccName -> OccName -> OccName
`T.append` (Char
s Char -> OccName -> OccName
`T.cons` String -> OccName
T.pack (Integer -> String
forall a. Show a => a -> String
show (Integer
maxNInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
n)))
        ,"_co_"
        ,"el"
        ,"rest"
        ]
        (Type
resTyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:Type
restTyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
idTys)

    elNVar :: Term
elNVar    = Var Term -> Term
Var Var Term
elNId
    pat :: Pat
pat       = DataCon -> [TyVar] -> [Var Term] -> Pat
DataPat DataCon
consCon [TyVar
mTV] [Var Term
co,Var Term
el,Var Term
rest]
    lhs :: Term
lhs       = Term -> Type -> [Alt] -> Term
Case Term
e Type
resTy  [(Pat
pat,Var Term -> Term
Var Var Term
el)]
    rhs :: Term
rhs       = Term -> Type -> [Alt] -> Term
Case Term
e Type
restTy [(Pat
pat,Var Term -> Term
Var Var Term
rest)]

    (uniqs3 :: (Supply, InScopeSet)
uniqs3,restVs :: [(Term, [LetBinding])]
restVs) = Integer
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), [(Term, [LetBinding])])
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) (Supply, InScopeSet)
uniqs2 (Var Term -> Term
Var Var Term
restNId)

-- | Create let-bindings with case-statements that select elements out of a
-- tree. Returns both the variables to which element-selections are bound
-- and the let-bindings
extractTElems
  :: Supply
  -- ^ Unique supply
  -> InScopeSet
  -- ^ (Superset of) in scope variables
  -> DataCon
  -- ^ The 'LR' constructor
  -> DataCon
  -- ^ The 'BR' constructor
  -> Type
  -- ^ The element type
  -> Char
  -- ^ Char to append to the bound variable names
  -> Integer
  -- ^ Depth of the tree
  -> Term
  -- ^ The tree
  -> (Supply,([Term],[LetBinding]))
extractTElems :: Supply
-> InScopeSet
-> DataCon
-> DataCon
-> Type
-> Char
-> Integer
-> Term
-> (Supply, ([Term], [LetBinding]))
extractTElems supply :: Supply
supply inScope :: InScopeSet
inScope lrCon :: DataCon
lrCon brCon :: DataCon
brCon resTy :: Type
resTy s :: Char
s maxN :: Integer
maxN tree :: Term
tree =
  ((Supply, InScopeSet) -> Supply)
-> ((Supply, InScopeSet), ([Term], [LetBinding]))
-> (Supply, ([Term], [LetBinding]))
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Supply, InScopeSet) -> Supply
forall a b. (a, b) -> a
fst (Integer
-> [Int]
-> [Int]
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), ([Term], [LetBinding]))
go Integer
maxN [0..(2Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^(Integer
maxNInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+1))Int -> Int -> Int
forall a. Num a => a -> a -> a
-2] [0..(2Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
maxN Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1)] (Supply
supply,InScopeSet
inScope) Term
tree)
 where
  go :: Integer
     -> [Int]
     -> [Int]
     -> (Supply,InScopeSet)
     -> Term
     -> ((Supply,InScopeSet),([Term],[LetBinding]))
  go :: Integer
-> [Int]
-> [Int]
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), ([Term], [LetBinding]))
go 0 _ ks :: [Int]
ks uniqs0 :: (Supply, InScopeSet)
uniqs0 e :: Term
e = ((Supply, InScopeSet)
uniqs1,([Term
elNVar],[(Var Term
elNId, Term
rhs)]))
   where
    tys :: [Type]
tys          = [LitTy -> Type
LitTy (Integer -> LitTy
NumTy 0),Type
resTy]
    (Just idTys :: [Type]
idTys) = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
lrCon [Type]
tys

    (uniqs1 :: (Supply, InScopeSet)
uniqs1,[elNId :: Var Term
elNId,co :: Var Term
co,el :: Var Term
el]) =
      ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var Term))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Var Term])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var Term)
mkUniqSystemId (Supply, InScopeSet)
uniqs0 ([(OccName, Type)] -> ((Supply, InScopeSet), [Var Term]))
-> [(OccName, Type)] -> ((Supply, InScopeSet), [Var Term])
forall a b. (a -> b) -> a -> b
$ [OccName] -> [Type] -> [(OccName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip
        [ "el" OccName -> OccName -> OccName
`T.append` (Char
s Char -> OccName -> OccName
`T.cons` String -> OccName
T.pack (Int -> String
forall a. Show a => a -> String
show ([Int] -> Int
forall a. [a] -> a
head [Int]
ks)))
        , "_co_"
        , "el"
        ]
        (Type
resTyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
idTys)
    elNVar :: Term
elNVar = Var Term -> Term
Var Var Term
elNId
    pat :: Pat
pat    = DataCon -> [TyVar] -> [Var Term] -> Pat
DataPat DataCon
lrCon [] [Var Term
co,Var Term
el]
    rhs :: Term
rhs    = Term -> Type -> [Alt] -> Term
Case Term
e Type
resTy [(Pat
pat,Var Term -> Term
Var Var Term
el)]

  go n :: Integer
n bs :: [Int]
bs ks :: [Int]
ks uniqs0 :: (Supply, InScopeSet)
uniqs0 e :: Term
e =
    ((Supply, InScopeSet)
uniqs4
    ,([Term]
lVars [Term] -> [Term] -> [Term]
forall a. [a] -> [a] -> [a]
++ [Term]
rVars,(Var Term
ltNId, Term
ltRhs)LetBinding -> [LetBinding] -> [LetBinding]
forall a. a -> [a] -> [a]
:
                     (Var Term
rtNId, Term
rtRhs)LetBinding -> [LetBinding] -> [LetBinding]
forall a. a -> [a] -> [a]
:
                     ([LetBinding]
lBinds [LetBinding] -> [LetBinding] -> [LetBinding]
forall a. [a] -> [a] -> [a]
++ [LetBinding]
rBinds)))
   where
    tys :: [Type]
tys = [LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n),Type
resTy,LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1))]
    (Just idTys :: [Type]
idTys) = DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
brCon [Type]
tys

    (uniqs1 :: (Supply, InScopeSet)
uniqs1,mTV :: TyVar
mTV) = (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar (Supply, InScopeSet)
uniqs0 ("m",Type
typeNatKind)
    (b0 :: Int
b0:bL :: [Int]
bL,b1 :: Int
b1:bR :: [Int]
bR) = Int -> [Int] -> ([Int], [Int])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
bs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` 2) [Int]
bs
    brTy :: Type
brTy = [Type] -> Type
forall a. [a] -> a
last [Type]
idTys
    (uniqs2 :: (Supply, InScopeSet)
uniqs2,[ltNId :: Var Term
ltNId,rtNId :: Var Term
rtNId,co :: Var Term
co,lt :: Var Term
lt,rt :: Var Term
rt]) =
      ((Supply, InScopeSet)
 -> (OccName, Type) -> ((Supply, InScopeSet), Var Term))
-> (Supply, InScopeSet)
-> [(OccName, Type)]
-> ((Supply, InScopeSet), [Var Term])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var Term)
mkUniqSystemId (Supply, InScopeSet)
uniqs1 ([(OccName, Type)] -> ((Supply, InScopeSet), [Var Term]))
-> [(OccName, Type)] -> ((Supply, InScopeSet), [Var Term])
forall a b. (a -> b) -> a -> b
$ [OccName] -> [Type] -> [(OccName, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip
        ["lt" OccName -> OccName -> OccName
`T.append` (Char
s Char -> OccName -> OccName
`T.cons` String -> OccName
T.pack (Int -> String
forall a. Show a => a -> String
show Int
b0))
        ,"rt" OccName -> OccName -> OccName
`T.append` (Char
s Char -> OccName -> OccName
`T.cons` String -> OccName
T.pack (Int -> String
forall a. Show a => a -> String
show Int
b1))
        ,"_co_"
        ,"lt"
        ,"rt"
        ]
        (Type
brTyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:Type
brTyType -> [Type] -> [Type]
forall a. a -> [a] -> [a]
:[Type]
idTys)
    ltVar :: Term
ltVar = Var Term -> Term
Var Var Term
ltNId
    rtVar :: Term
rtVar = Var Term -> Term
Var Var Term
rtNId
    pat :: Pat
pat   = DataCon -> [TyVar] -> [Var Term] -> Pat
DataPat DataCon
brCon [TyVar
mTV] [Var Term
co,Var Term
lt,Var Term
rt]
    ltRhs :: Term
ltRhs = Term -> Type -> [Alt] -> Term
Case Term
e Type
brTy [(Pat
pat,Var Term -> Term
Var Var Term
lt)]
    rtRhs :: Term
rtRhs = Term -> Type -> [Alt] -> Term
Case Term
e Type
brTy [(Pat
pat,Var Term -> Term
Var Var Term
rt)]

    (kL :: [Int]
kL,kR :: [Int]
kR) = Int -> [Int] -> ([Int], [Int])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ks Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` 2) [Int]
ks
    (uniqs3 :: (Supply, InScopeSet)
uniqs3,(lVars :: [Term]
lVars,lBinds :: [LetBinding]
lBinds)) = Integer
-> [Int]
-> [Int]
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), ([Term], [LetBinding]))
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Int]
bL [Int]
kL (Supply, InScopeSet)
uniqs2 Term
ltVar
    (uniqs4 :: (Supply, InScopeSet)
uniqs4,(rVars :: [Term]
rVars,rBinds :: [LetBinding]
rBinds)) = Integer
-> [Int]
-> [Int]
-> (Supply, InScopeSet)
-> Term
-> ((Supply, InScopeSet), ([Term], [LetBinding]))
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Int]
bR [Int]
kR (Supply, InScopeSet)
uniqs3 Term
rtVar

-- | Create a vector of supplied elements
mkRTree :: DataCon -- ^ The LR constructor
        -> DataCon -- ^ The BR constructor
        -> Type    -- ^ Element type
        -> Integer -- ^ Depth of the tree
        -> [Term]  -- ^ Elements to put in the tree
        -> Term
mkRTree :: DataCon -> DataCon -> Type -> Integer -> [Term] -> Term
mkRTree lrCon :: DataCon
lrCon brCon :: DataCon
brCon resTy :: Type
resTy = Integer -> [Term] -> Term
go
  where
    go :: Integer -> [Term] -> Term
go _ [x :: Term
x] = Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
lrCon) [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy 0))
                                    ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                                    ,Term -> Either Term Type
forall a b. a -> Either a b
Left  (Type -> Term
primCo Type
lrCoTy)
                                    ,Term -> Either Term Type
forall a b. a -> Either a b
Left  Term
x
                                    ]

    go n :: Integer
n xs :: [Term]
xs =
      let (xsL :: [Term]
xsL,xsR :: [Term]
xsR) = Int -> [Term] -> ([Term], [Term])
forall a. Int -> [a] -> ([a], [a])
splitAt ([Term] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Term]
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` 2) [Term]
xs
      in  Term -> [Either Term Type] -> Term
mkApps (DataCon -> Term
Data DataCon
brCon) [Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                              ,Type -> Either Term Type
forall a b. b -> Either a b
Right Type
resTy
                              ,Type -> Either Term Type
forall a b. b -> Either a b
Right (LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))
                              ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Type -> Term
primCo (Integer -> Type
brCoTy Integer
n))
                              ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> [Term] -> Term
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Term]
xsL)
                              ,Term -> Either Term Type
forall a b. a -> Either a b
Left (Integer -> [Term] -> Term
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1) [Term]
xsR)]

    lrCoTy :: Type
lrCoTy   = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Type] -> [Type]) -> Maybe [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$! DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
lrCon  [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy 0))
                                                         ,Type
resTy])
    brCoTy :: Integer -> Type
brCoTy n :: Integer
n = [Type] -> Type
forall a. [a] -> a
head (Maybe [Type] -> [Type]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [Type] -> [Type]) -> Maybe [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$! DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys DataCon
brCon
                                                   [(LitTy -> Type
LitTy (Integer -> LitTy
NumTy Integer
n))
                                                   ,Type
resTy
                                                   ,(LitTy -> Type
LitTy (Integer -> LitTy
NumTy (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-1)))])

-- | Determine whether a type is isomorphic to "Clash.Signal.Internal.Signal"
--
-- It is i.e.:
--
--   * Signal clk a
--   * (Signal clk a, Signal clk b)
--   * Vec n (Signal clk a)
--   * data Wrap = W (Signal clk' Int)
--   * etc.
--
-- This also includes BiSignals, i.e.:
--
--   * BiSignalIn High System Int
--   * etc.
--
isSignalType :: TyConMap -> Type -> Bool
isSignalType :: TyConMap -> Type -> Bool
isSignalType tcm :: TyConMap
tcm ty :: Type
ty = HashSet TyConName -> Type -> Bool
go HashSet TyConName
forall a. HashSet a
HashSet.empty Type
ty
  where
    go :: HashSet TyConName -> Type -> Bool
go tcSeen :: HashSet TyConName
tcSeen (Type -> TypeView
tyView -> TyConApp tcNm :: TyConName
tcNm args :: [Type]
args) = case TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm of
      "Clash.Signal.Internal.Signal"      -> Bool
True
      "Clash.Signal.BiSignal.BiSignalIn"  -> Bool
True
      "Clash.Signal.Internal.BiSignalOut" -> Bool
True
      _ | TyConName
tcNm TyConName -> HashSet TyConName -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`HashSet.member` HashSet TyConName
tcSeen    -> Bool
False -- Do not follow rec types
        | Bool
otherwise -> case TyConName -> TyConMap -> Maybe TyCon
forall a b. Uniquable a => a -> UniqMap b -> Maybe b
lookupUniqMap TyConName
tcNm TyConMap
tcm of
            Just tc :: TyCon
tc -> let dcs :: [DataCon]
dcs         = TyCon -> [DataCon]
tyConDataCons TyCon
tc
                           dcInsArgTys :: [Type]
dcInsArgTys = [[Type]] -> [Type]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
                                       ([[Type]] -> [Type]) -> [[Type]] -> [Type]
forall a b. (a -> b) -> a -> b
$ (DataCon -> Maybe [Type]) -> [DataCon] -> [[Type]]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DataCon -> [Type] -> Maybe [Type]
`dataConInstArgTys` [Type]
args) [DataCon]
dcs
                           tcSeen' :: HashSet TyConName
tcSeen'     = TyConName -> HashSet TyConName -> HashSet TyConName
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert TyConName
tcNm HashSet TyConName
tcSeen
                       in  (Type -> Bool) -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (HashSet TyConName -> Type -> Bool
go HashSet TyConName
tcSeen') [Type]
dcInsArgTys
            Nothing -> Bool -> String -> Bool -> Bool
forall a. Bool -> String -> a -> a
traceIf Bool
True ($(curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "isSignalType: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ TyConName -> String
forall a. Show a => a -> String
show TyConName
tcNm
                                     String -> ShowS
forall a. [a] -> [a] -> [a]
++ " not found.") Bool
False

    go _ _ = Bool
False

isClockOrReset
  :: TyConMap
  -> Type
  -> Bool
isClockOrReset :: TyConMap -> Type -> Bool
isClockOrReset m :: TyConMap
m (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m -> Just ty :: Type
ty)    = TyConMap -> Type -> Bool
isClockOrReset TyConMap
m Type
ty
isClockOrReset _ (Type -> TypeView
tyView -> TyConApp tcNm :: TyConName
tcNm _) = case TyConName -> OccName
forall a. Name a -> OccName
nameOcc TyConName
tcNm of
  "Clash.Signal.Internal.Clock" -> Bool
True
  "Clash.Signal.Internal.Reset" -> Bool
True
  _ -> Bool
False
isClockOrReset _ _ = Bool
False

tyNatSize :: TyConMap
          -> Type
          -> Except String Integer
tyNatSize :: TyConMap -> Type -> Except String Integer
tyNatSize m :: TyConMap
m (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m -> Just ty :: Type
ty) = TyConMap -> Type -> Except String Integer
tyNatSize TyConMap
m Type
ty
tyNatSize _ (LitTy (NumTy i :: Integer
i))        = Integer -> Except String Integer
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
i
tyNatSize _ ty :: Type
ty = String -> Except String Integer
forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
throwE (String -> Except String Integer)
-> String -> Except String Integer
forall a b. (a -> b) -> a -> b
$ $(curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "Cannot reduce to an integer:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall p. PrettyPrec p => p -> String
showPpr Type
ty


mkUniqSystemTyVar
  :: (Supply, InScopeSet)
  -> (OccName, Kind)
  -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar :: (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), TyVar)
mkUniqSystemTyVar (supply :: Supply
supply,inScope :: InScopeSet
inScope) (nm :: OccName
nm, ki :: Type
ki) =
  ((Supply
supply',InScopeSet -> TyVar -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope TyVar
v'), TyVar
v')
 where
  (u :: Int
u,supply' :: Supply
supply') = Supply -> (Int, Supply)
freshId Supply
supply
  v :: TyVar
v           = Type -> TyName -> TyVar
mkTyVar Type
ki (OccName -> Int -> TyName
forall a. OccName -> Int -> Name a
mkUnsafeSystemName OccName
nm Int
u)
  v' :: TyVar
v'          = InScopeSet -> TyVar -> TyVar
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope TyVar
v

mkUniqSystemId
  :: (Supply, InScopeSet)
  -> (OccName, Type)
  -> ((Supply,InScopeSet), Id)
mkUniqSystemId :: (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var Term)
mkUniqSystemId (supply :: Supply
supply,inScope :: InScopeSet
inScope) (nm :: OccName
nm, ty :: Type
ty) =
  ((Supply
supply',InScopeSet -> Var Term -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope Var Term
v'), Var Term
v')
 where
  (u :: Int
u,supply' :: Supply
supply') = Supply -> (Int, Supply)
freshId Supply
supply
  v :: Var Term
v           = Type -> TmName -> Var Term
mkLocalId Type
ty (OccName -> Int -> TmName
forall a. OccName -> Int -> Name a
mkUnsafeSystemName OccName
nm Int
u)
  v' :: Var Term
v'          = InScopeSet -> Var Term -> Var Term
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope Var Term
v

mkUniqInternalId
  :: (Supply, InScopeSet)
  -> (OccName, Type)
  -> ((Supply,InScopeSet), Id)
mkUniqInternalId :: (Supply, InScopeSet)
-> (OccName, Type) -> ((Supply, InScopeSet), Var Term)
mkUniqInternalId (supply :: Supply
supply,inScope :: InScopeSet
inScope) (nm :: OccName
nm, ty :: Type
ty) =
  ((Supply
supply',InScopeSet -> Var Term -> InScopeSet
forall a. InScopeSet -> Var a -> InScopeSet
extendInScopeSet InScopeSet
inScope Var Term
v'), Var Term
v')
 where
  (u :: Int
u,supply' :: Supply
supply') = Supply -> (Int, Supply)
freshId Supply
supply
  v :: Var Term
v           = Type -> TmName -> Var Term
mkLocalId Type
ty (OccName -> Int -> TmName
forall a. OccName -> Int -> Name a
mkUnsafeInternalName OccName
nm Int
u)
  v' :: Var Term
v'          = InScopeSet -> Var Term -> Var Term
forall a. InScopeSet -> Var a -> Var a
uniqAway InScopeSet
inScope Var Term
v


-- | Same as @dataConInstArgTys@, but it tries to compute existentials too,
-- hence the extra argument @TyConMap@. WARNING: It will return the types
-- of non-existentials only
dataConInstArgTysE
  :: HasCallStack
  => InScopeSet
  -> TyConMap
  -> DataCon
  -> [Type]
  -> Maybe [Type]
dataConInstArgTysE :: InScopeSet -> TyConMap -> DataCon -> [Type] -> Maybe [Type]
dataConInstArgTysE is0 :: InScopeSet
is0 tcm :: TyConMap
tcm (MkData { [Type]
dcArgTys :: [Type]
dcArgTys :: DataCon -> [Type]
dcArgTys, [TyVar]
dcExtTyVars :: [TyVar]
dcExtTyVars :: DataCon -> [TyVar]
dcExtTyVars, [TyVar]
dcUnivTyVars :: [TyVar]
dcUnivTyVars :: DataCon -> [TyVar]
dcUnivTyVars }) inst_tys :: [Type]
inst_tys = do
  -- TODO: Check if all existentials were solved (they should be, or the wouldn't have
  -- TODO: been solved in the caseElemExistentials transformation)
  let is1 :: InScopeSet
is1   = InScopeSet -> [TyVar] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 [TyVar]
dcExtTyVars
      is2 :: InScopeSet
is2   = InScopeSet -> InScopeSet -> InScopeSet
unionInScope InScopeSet
is1 (VarSet -> InScopeSet
mkInScopeSet ([Type] -> VarSet
forall (f :: * -> *). Foldable f => f Type -> VarSet
tyFVsOfTypes [Type]
inst_tys))
      subst :: Subst
subst = Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (InScopeSet -> Subst
mkSubst InScopeSet
is2) ([TyVar] -> [Type] -> [(TyVar, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
dcUnivTyVars [Type]
inst_tys)
  [TyVar] -> [Type] -> Maybe [Type]
go
    (HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substGlobalsInExistentials InScopeSet
is0 [TyVar]
dcExtTyVars ([TyVar] -> [Type] -> [(TyVar, Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
dcUnivTyVars [Type]
inst_tys))
    ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst) [Type]
dcArgTys)

 where
  go
    :: [TyVar]
    -- ^ Existentials
    -> [Type]
    -- ^ Type arguments
    -> Maybe [Type]
    -- ^ Maybe ([type of non-existential])
  go :: [TyVar] -> [Type] -> Maybe [Type]
go exts0 :: [TyVar]
exts0 args0 :: [Type]
args0 =
    let eqs :: [(Type, Type)]
eqs = [Maybe (Type, Type)] -> [(Type, Type)]
forall a. [Maybe a] -> [a]
catMaybes ((Type -> Maybe (Type, Type)) -> [Type] -> [Maybe (Type, Type)]
forall a b. (a -> b) -> [a] -> [b]
map (TyConMap -> Type -> Maybe (Type, Type)
typeEq TyConMap
tcm) [Type]
args0) in
    case TyConMap -> [(Type, Type)] -> [(TyVar, Type)]
solveNonAbsurds TyConMap
tcm [(Type, Type)]
eqs of
      [] ->
        [Type] -> Maybe [Type]
forall a. a -> Maybe a
Just [Type]
args0
      sols :: [(TyVar, Type)]
sols ->
        [TyVar] -> [Type] -> Maybe [Type]
go [TyVar]
exts1 [Type]
args1
        where
          exts1 :: [TyVar]
exts1 = HasCallStack => InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
InScopeSet -> [TyVar] -> [(TyVar, Type)] -> [TyVar]
substInExistentialsList InScopeSet
is0 [TyVar]
exts0 [(TyVar, Type)]
sols
          is2 :: InScopeSet
is2   = InScopeSet -> [TyVar] -> InScopeSet
forall a. InScopeSet -> [Var a] -> InScopeSet
extendInScopeSetList InScopeSet
is0 [TyVar]
exts1
          subst :: Subst
subst = Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (InScopeSet -> Subst
mkSubst InScopeSet
is2) [(TyVar, Type)]
sols
          args1 :: [Type]
args1 = (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst) [Type]
args0


-- | Given a DataCon and a list of types, the type variables of the DataCon
-- type are substituted for the list of types. The argument types are returned.
--
-- The list of types should be equal to the number of type variables, otherwise
-- @Nothing@ is returned.
dataConInstArgTys :: DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys :: DataCon -> [Type] -> Maybe [Type]
dataConInstArgTys (MkData { [Type]
dcArgTys :: [Type]
dcArgTys :: DataCon -> [Type]
dcArgTys, [TyVar]
dcUnivTyVars :: [TyVar]
dcUnivTyVars :: DataCon -> [TyVar]
dcUnivTyVars, [TyVar]
dcExtTyVars :: [TyVar]
dcExtTyVars :: DataCon -> [TyVar]
dcExtTyVars }) inst_tys :: [Type]
inst_tys =
  -- TODO: Check if inst_tys do not contain any free variables on call sites. If
  -- TODO: they do, this function is unsafe to use.
  let tyvars :: [TyVar]
tyvars = [TyVar]
dcUnivTyVars [TyVar] -> [TyVar] -> [TyVar]
forall a. [a] -> [a] -> [a]
++ [TyVar]
dcExtTyVars in
  if [TyVar] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVar]
tyvars Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
inst_tys then
    [Type] -> Maybe [Type]
forall a. a -> Maybe a
Just ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => [TyVar] -> [Type] -> Type -> Type
[TyVar] -> [Type] -> Type -> Type
substTyWith [TyVar]
tyvars [Type]
inst_tys) [Type]
dcArgTys)
  else
    Maybe [Type]
forall a. Maybe a
Nothing

-- | Make a coercion
primCo
  :: Type
  -> Term
primCo :: Type -> Term
primCo ty :: Type
ty = OccName -> PrimInfo -> Term
Prim "_CO_" (Type -> WorkInfo -> PrimInfo
PrimInfo Type
ty WorkInfo
WorkNever)

-- | Make an undefined term
undefinedTm
  :: Type
  -> Term
undefinedTm :: Type -> Term
undefinedTm = Term -> Type -> Term
TyApp (OccName -> PrimInfo -> Term
Prim "Clash.Transformations.undefined" (Type -> WorkInfo -> PrimInfo
PrimInfo Type
undefinedTy WorkInfo
WorkNever))

substArgTys
  :: DataCon
  -> [Type]
  -> [Type]
substArgTys :: DataCon -> [Type] -> [Type]
substArgTys dc :: DataCon
dc args :: [Type]
args =
  let univTVs :: [TyVar]
univTVs = DataCon -> [TyVar]
dcUnivTyVars DataCon
dc
      extTVs :: [TyVar]
extTVs  = DataCon -> [TyVar]
dcExtTyVars DataCon
dc
      argsFVs :: VarSet
argsFVs = (VarSet -> VarSet -> VarSet) -> VarSet -> [VarSet] -> VarSet
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' VarSet -> VarSet -> VarSet
unionVarSet VarSet
emptyVarSet
                  ((Type -> VarSet) -> [Type] -> [VarSet]
forall a b. (a -> b) -> [a] -> [b]
map (Getting VarSet Type TyVar -> (TyVar -> VarSet) -> Type -> VarSet
forall r s a. Getting r s a -> (a -> r) -> s -> r
Lens.foldMapOf Getting VarSet Type TyVar
Fold Type TyVar
typeFreeVars TyVar -> VarSet
forall a. Var a -> VarSet
unitVarSet) [Type]
args)
      is :: InScopeSet
is      = VarSet -> InScopeSet
mkInScopeSet (VarSet
argsFVs VarSet -> VarSet -> VarSet
`unionVarSet` [TyVar] -> VarSet
forall a. [Var a] -> VarSet
mkVarSet [TyVar]
extTVs)
      -- See Note [The substitution invariant]
      subst :: Subst
subst   = Subst -> [(TyVar, Type)] -> Subst
extendTvSubstList (InScopeSet -> Subst
mkSubst InScopeSet
is) ([TyVar]
univTVs [TyVar] -> [Type] -> [(TyVar, Type)]
forall a b. [a] -> [b] -> [(a, b)]
`zipEqual` [Type]
args)
  in  (Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (HasCallStack => Subst -> Type -> Type
Subst -> Type -> Type
substTy Subst
subst) (DataCon -> [Type]
dcArgTys DataCon
dc)

stripTicks :: Term -> Term
stripTicks :: Term -> Term
stripTicks (Tick _ e :: Term
e) = Term -> Term
stripTicks Term
e
stripTicks e :: Term
e = Term
e

-- | Try to reduce an arbitrary type to a Symbol, and subsequently extract its
-- String representation
tySym
  :: TyConMap
  -> Type
  -> Except String String
tySym :: TyConMap -> Type -> Except String String
tySym m :: TyConMap
m (TyConMap -> Type -> Maybe Type
coreView1 TyConMap
m -> Just ty :: Type
ty) = TyConMap -> Type -> Except String String
tySym TyConMap
m Type
ty
tySym _ (LitTy (SymTy s :: String
s))        = String -> Except String String
forall (m :: * -> *) a. Monad m => a -> m a
return String
s
tySym _ ty :: Type
ty = String -> Except String String
forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
throwE (String -> Except String String) -> String -> Except String String
forall a b. (a -> b) -> a -> b
$ $(curLoc) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "Cannot reduce to a string:\n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall p. PrettyPrec p => p -> String
showPpr Type
ty