-- | The notion of scope used during term checking, and other small
-- definitions shared between the unsized term checker
-- ("Language.Futhark.TypeChecker.Terms.Unsized") and the sized one
-- ("Language.Futhark.TypeChecker.Terms"). The two checkers represent
-- sizes differently, so the definitions here are parameterised over
-- the size representation, and the functions take a function for
-- converting the sizes of module-level types.
module Language.Futhark.TypeChecker.Terms.Scope
  ( ValBinding (..),
    TermScope (..),
    Inferred (..),
    envToTermScope,
    envToTermScopeNoVals,
    lookupOuterVal,
    initialTermScope,
    lookupQualNameEnv,
    lookupQualNameEnvMaybe,
    typeParamIdent,
  )
where

import Data.Map.Strict qualified as M
import Data.Maybe
import Language.Futhark
import Language.Futhark.TypeChecker.Monad hiding (BoundV)
import Language.Futhark.TypeChecker.Monad qualified as TypeM

-- | What a bound value stands for. Note that although modules are in the same
-- name space, they are not value bindings.
data ValBinding dim
  = BoundV [TypeParam] (TypeBase dim NoMode)
  | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType)
  | EqualityF
  | -- | A recursive function with no declared return type, currently being
    -- checked and in scope of its own body. Each occurrence adopts the type
    -- recorded for it by the unsized type checker, with fresh (unrelated) sizes.
    -- See Note [Checking recursive functions] in
    -- Language.Futhark.TypeChecker.Terms.
    RecursiveV
  deriving (Int -> ValBinding dim -> ShowS
[ValBinding dim] -> ShowS
ValBinding dim -> String
(Int -> ValBinding dim -> ShowS)
-> (ValBinding dim -> String)
-> ([ValBinding dim] -> ShowS)
-> Show (ValBinding dim)
forall dim. Show dim => Int -> ValBinding dim -> ShowS
forall dim. Show dim => [ValBinding dim] -> ShowS
forall dim. Show dim => ValBinding dim -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall dim. Show dim => Int -> ValBinding dim -> ShowS
showsPrec :: Int -> ValBinding dim -> ShowS
$cshow :: forall dim. Show dim => ValBinding dim -> String
show :: ValBinding dim -> String
$cshowList :: forall dim. Show dim => [ValBinding dim] -> ShowS
showList :: [ValBinding dim] -> ShowS
Show)

-- | The lexical scope information.
data TermScope dim = TermScope
  { forall dim. TermScope dim -> Map VName (ValBinding dim)
scopeVtable :: M.Map VName (ValBinding dim),
    forall dim. TermScope dim -> Map VName TypeBinding
scopeTypeTable :: M.Map VName TypeBinding,
    forall dim. TermScope dim -> Map VName Mod
scopeModTable :: M.Map VName Mod
  }
  deriving (Int -> TermScope dim -> ShowS
[TermScope dim] -> ShowS
TermScope dim -> String
(Int -> TermScope dim -> ShowS)
-> (TermScope dim -> String)
-> ([TermScope dim] -> ShowS)
-> Show (TermScope dim)
forall dim. Show dim => Int -> TermScope dim -> ShowS
forall dim. Show dim => [TermScope dim] -> ShowS
forall dim. Show dim => TermScope dim -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall dim. Show dim => Int -> TermScope dim -> ShowS
showsPrec :: Int -> TermScope dim -> ShowS
$cshow :: forall dim. Show dim => TermScope dim -> String
show :: TermScope dim -> String
$cshowList :: forall dim. Show dim => [TermScope dim] -> ShowS
showList :: [TermScope dim] -> ShowS
Show)

instance Semigroup (TermScope dim) where
  TermScope Map VName (ValBinding dim)
vt1 Map VName TypeBinding
tt1 Map VName Mod
mt1 <> :: TermScope dim -> TermScope dim -> TermScope dim
<> TermScope Map VName (ValBinding dim)
vt2 Map VName TypeBinding
tt2 Map VName Mod
mt2 =
    Map VName (ValBinding dim)
-> Map VName TypeBinding -> Map VName Mod -> TermScope dim
forall dim.
Map VName (ValBinding dim)
-> Map VName TypeBinding -> Map VName Mod -> TermScope dim
TermScope (Map VName (ValBinding dim)
vt2 Map VName (ValBinding dim)
-> Map VName (ValBinding dim) -> Map VName (ValBinding dim)
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName (ValBinding dim)
vt1) (Map VName TypeBinding
tt2 Map VName TypeBinding
-> Map VName TypeBinding -> Map VName TypeBinding
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName TypeBinding
tt1) (Map VName Mod
mt1 Map VName Mod -> Map VName Mod -> Map VName Mod
forall k a. Ord k => Map k a -> Map k a -> Map k a
`M.union` Map VName Mod
mt2)

-- | During type checking a pattern, we might find an explicit ascription. These
-- contain complete type information (although they must of course still be
-- checked against what remains of the pattern).
data Inferred t
  = NoneInferred
  | Ascribed t
  deriving (Int -> Inferred t -> ShowS
[Inferred t] -> ShowS
Inferred t -> String
(Int -> Inferred t -> ShowS)
-> (Inferred t -> String)
-> ([Inferred t] -> ShowS)
-> Show (Inferred t)
forall t. Show t => Int -> Inferred t -> ShowS
forall t. Show t => [Inferred t] -> ShowS
forall t. Show t => Inferred t -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall t. Show t => Int -> Inferred t -> ShowS
showsPrec :: Int -> Inferred t -> ShowS
$cshow :: forall t. Show t => Inferred t -> String
show :: Inferred t -> String
$cshowList :: forall t. Show t => [Inferred t] -> ShowS
showList :: [Inferred t] -> ShowS
Show)

instance Functor Inferred where
  fmap :: forall a b. (a -> b) -> Inferred a -> Inferred b
fmap a -> b
_ Inferred a
NoneInferred = Inferred b
forall t. Inferred t
NoneInferred
  fmap a -> b
f (Ascribed a
t) = b -> Inferred b
forall t. t -> Inferred t
Ascribed (a -> b
f a
t)

-- | Create a scope from a module-level environment.
envToTermScope ::
  (StructType -> TypeBase dim NoMode) ->
  Env ->
  TermScope dim
envToTermScope :: forall dim.
(StructType -> TypeBase dim NoMode) -> Env -> TermScope dim
envToTermScope StructType -> TypeBase dim NoMode
onType Env
env =
  TermScope
    { scopeVtable :: Map VName (ValBinding dim)
scopeVtable = Map VName (ValBinding dim)
vtable,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = Env -> Map VName TypeBinding
envTypeTable Env
env,
      scopeModTable :: Map VName Mod
scopeModTable = Env -> Map VName Mod
envModTable Env
env
    }
  where
    vtable :: Map VName (ValBinding dim)
vtable = (BoundV -> ValBinding dim)
-> Map VName BoundV -> Map VName (ValBinding dim)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map BoundV -> ValBinding dim
valBinding (Map VName BoundV -> Map VName (ValBinding dim))
-> Map VName BoundV -> Map VName (ValBinding dim)
forall a b. (a -> b) -> a -> b
$ Env -> Map VName BoundV
envVtable Env
env
    valBinding :: BoundV -> ValBinding dim
valBinding (TypeM.BoundV [TypeParam]
tps StructType
v) = [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
forall dim. [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParam]
tps (TypeBase dim NoMode -> ValBinding dim)
-> TypeBase dim NoMode -> ValBinding dim
forall a b. (a -> b) -> a -> b
$ StructType -> TypeBase dim NoMode
onType StructType
v

-- | Like 'envToTermScope', but omits the (potentially very large)
-- value table. Its bindings are instead looked up on demand with
-- 'lookupOuterVal'. This avoids transforming the entire value table of
-- the outer environment - which for a top-level definition includes
-- the whole prelude - once for every binding we check.
envToTermScopeNoVals :: Env -> TermScope dim
envToTermScopeNoVals :: forall dim. Env -> TermScope dim
envToTermScopeNoVals Env
env =
  TermScope
    { scopeVtable :: Map VName (ValBinding dim)
scopeVtable = Map VName (ValBinding dim)
forall a. Monoid a => a
mempty,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = Env -> Map VName TypeBinding
envTypeTable Env
env,
      scopeModTable :: Map VName Mod
scopeModTable = Env -> Map VName Mod
envModTable Env
env
    }

-- | Look up a single value binding in an 'Env' and convert it, using
-- the given size conversion. The fallback for names not found in the
-- (value-free) term scope built by 'envToTermScopeNoVals'.
lookupOuterVal ::
  (StructType -> TypeBase dim NoMode) ->
  Env ->
  VName ->
  Maybe (ValBinding dim)
lookupOuterVal :: forall dim.
(StructType -> TypeBase dim NoMode)
-> Env -> VName -> Maybe (ValBinding dim)
lookupOuterVal StructType -> TypeBase dim NoMode
onType Env
env VName
v =
  BoundV -> ValBinding dim
convert (BoundV -> ValBinding dim)
-> Maybe BoundV -> Maybe (ValBinding dim)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> VName -> Map VName BoundV -> Maybe BoundV
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
v (Env -> Map VName BoundV
envVtable Env
env)
  where
    convert :: BoundV -> ValBinding dim
convert (TypeM.BoundV [TypeParam]
tps StructType
t) = [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
forall dim. [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParam]
tps (TypeBase dim NoMode -> ValBinding dim)
-> TypeBase dim NoMode -> ValBinding dim
forall a b. (a -> b) -> a -> b
$ StructType -> TypeBase dim NoMode
onType StructType
t

-- | The initial scope, containing the intrinsics.
initialTermScope ::
  (StructType -> TypeBase dim NoMode) ->
  TermScope dim
initialTermScope :: forall dim. (StructType -> TypeBase dim NoMode) -> TermScope dim
initialTermScope StructType -> TypeBase dim NoMode
onType =
  TermScope
    { scopeVtable :: Map VName (ValBinding dim)
scopeVtable = Map VName (ValBinding dim)
initialVtable,
      scopeTypeTable :: Map VName TypeBinding
scopeTypeTable = Map VName TypeBinding
forall a. Monoid a => a
mempty,
      scopeModTable :: Map VName Mod
scopeModTable = Map VName Mod
forall a. Monoid a => a
mempty
    }
  where
    initialVtable :: Map VName (ValBinding dim)
initialVtable = [(VName, ValBinding dim)] -> Map VName (ValBinding dim)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(VName, ValBinding dim)] -> Map VName (ValBinding dim))
-> [(VName, ValBinding dim)] -> Map VName (ValBinding dim)
forall a b. (a -> b) -> a -> b
$ ((VName, Intrinsic) -> Maybe (VName, ValBinding dim))
-> [(VName, Intrinsic)] -> [(VName, ValBinding dim)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (VName, Intrinsic) -> Maybe (VName, ValBinding dim)
forall {a}. (a, Intrinsic) -> Maybe (a, ValBinding dim)
addIntrinsicF ([(VName, Intrinsic)] -> [(VName, ValBinding dim)])
-> [(VName, Intrinsic)] -> [(VName, ValBinding dim)]
forall a b. (a -> b) -> a -> b
$ Map VName Intrinsic -> [(VName, Intrinsic)]
forall k a. Map k a -> [(k, a)]
M.toList Map VName Intrinsic
intrinsics

    prim :: PrimType -> TypeBase dim o
prim = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> (PrimType -> ScalarTypeBase dim o) -> PrimType -> TypeBase dim o
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimType -> ScalarTypeBase dim o
forall dim o. PrimType -> ScalarTypeBase dim o
Prim
    arrow :: TypeBase dim NoMode -> RetTypeBase dim Freshness -> TypeBase dim o
arrow TypeBase dim NoMode
x RetTypeBase dim Freshness
y = ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
forall dim o.
o
-> PName
-> Diet
-> TypeBase dim NoMode
-> RetTypeBase dim Freshness
-> ScalarTypeBase dim o
Arrow o
forall a. Monoid a => a
mempty PName
Unnamed Diet
Observe TypeBase dim NoMode
x RetTypeBase dim Freshness
y

    addIntrinsicF :: (a, Intrinsic) -> Maybe (a, ValBinding dim)
addIntrinsicF (a
name, IntrinsicMonoFun [PrimType]
pts PrimType
t) =
      (a, ValBinding dim) -> Maybe (a, ValBinding dim)
forall a. a -> Maybe a
Just (a
name, [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
forall dim. [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
BoundV [] (TypeBase dim NoMode -> ValBinding dim)
-> TypeBase dim NoMode -> ValBinding dim
forall a b. (a -> b) -> a -> b
$ StructType -> TypeBase dim NoMode
onType (StructType -> TypeBase dim NoMode)
-> StructType -> TypeBase dim NoMode
forall a b. (a -> b) -> a -> b
$ StructType -> RetTypeBase Size Freshness -> StructType
forall {o} {dim}.
Monoid o =>
TypeBase dim NoMode -> RetTypeBase dim Freshness -> TypeBase dim o
arrow StructType
forall {dim} {o}. TypeBase dim o
pts' (RetTypeBase Size Freshness -> StructType)
-> RetTypeBase Size Freshness -> StructType
forall a b. (a -> b) -> a -> b
$ [VName] -> TypeBase Size Freshness -> RetTypeBase Size Freshness
forall dim o. [VName] -> TypeBase dim o -> RetTypeBase dim o
RetType [] (TypeBase Size Freshness -> RetTypeBase Size Freshness)
-> TypeBase Size Freshness -> RetTypeBase Size Freshness
forall a b. (a -> b) -> a -> b
$ PrimType -> TypeBase Size Freshness
forall {dim} {o}. PrimType -> TypeBase dim o
prim PrimType
t)
      where
        pts' :: TypeBase dim o
pts' = case [PrimType]
pts of
          [PrimType
pt] -> PrimType -> TypeBase dim o
forall {dim} {o}. PrimType -> TypeBase dim o
prim PrimType
pt
          [PrimType]
_ -> ScalarTypeBase dim o -> TypeBase dim o
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase dim o -> TypeBase dim o)
-> ScalarTypeBase dim o -> TypeBase dim o
forall a b. (a -> b) -> a -> b
$ [TypeBase dim o] -> ScalarTypeBase dim o
forall dim o. [TypeBase dim o] -> ScalarTypeBase dim o
tupleRecord ([TypeBase dim o] -> ScalarTypeBase dim o)
-> [TypeBase dim o] -> ScalarTypeBase dim o
forall a b. (a -> b) -> a -> b
$ (PrimType -> TypeBase dim o) -> [PrimType] -> [TypeBase dim o]
forall a b. (a -> b) -> [a] -> [b]
map PrimType -> TypeBase dim o
forall {dim} {o}. PrimType -> TypeBase dim o
prim [PrimType]
pts
    addIntrinsicF (a
name, IntrinsicOverloadedFun [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts) =
      (a, ValBinding dim) -> Maybe (a, ValBinding dim)
forall a. a -> Maybe a
Just (a
name, [PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding dim
forall dim.
[PrimType] -> [Maybe PrimType] -> Maybe PrimType -> ValBinding dim
OverloadedF [PrimType]
ts [Maybe PrimType]
pts Maybe PrimType
rts)
    addIntrinsicF (a
name, IntrinsicPolyFun [TypeParam]
tvs [ParamType]
pts RetTypeBase Size Freshness
rt) =
      (a, ValBinding dim) -> Maybe (a, ValBinding dim)
forall a. a -> Maybe a
Just
        ( a
name,
          [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
forall dim. [TypeParam] -> TypeBase dim NoMode -> ValBinding dim
BoundV [TypeParam]
tvs (TypeBase dim NoMode -> ValBinding dim)
-> TypeBase dim NoMode -> ValBinding dim
forall a b. (a -> b) -> a -> b
$ StructType -> TypeBase dim NoMode
onType (StructType -> TypeBase dim NoMode)
-> StructType -> TypeBase dim NoMode
forall a b. (a -> b) -> a -> b
$ [ParamType] -> RetTypeBase Size Freshness -> StructType
forall d.
[TypeBase d Diet] -> RetTypeBase d Freshness -> TypeBase d NoMode
foldFunType [ParamType]
pts RetTypeBase Size Freshness
rt
        )
    addIntrinsicF (a
name, Intrinsic
IntrinsicEquality) =
      (a, ValBinding dim) -> Maybe (a, ValBinding dim)
forall a. a -> Maybe a
Just (a
name, ValBinding dim
forall dim. ValBinding dim
EqualityF)
    addIntrinsicF (a, Intrinsic)
_ = Maybe (a, ValBinding dim)
forall a. Maybe a
Nothing

-- | Find the scope corresponding to the qualifiers of the given
-- name. Fails with 'error' if the qualifiers do not name a module,
-- as this means the program should not have made it through earlier
-- checks.
lookupQualNameEnv ::
  (StructType -> TypeBase dim NoMode) ->
  TermScope dim ->
  QualName VName ->
  TermScope dim
lookupQualNameEnv :: forall dim.
(StructType -> TypeBase dim NoMode)
-> TermScope dim -> QualName VName -> TermScope dim
lookupQualNameEnv StructType -> TypeBase dim NoMode
onType TermScope dim
scope QualName VName
qn =
  TermScope dim -> Maybe (TermScope dim) -> TermScope dim
forall a. a -> Maybe a -> a
fromMaybe (String -> TermScope dim
forall a. HasCallStack => String -> a
error (String -> TermScope dim) -> String -> TermScope dim
forall a b. (a -> b) -> a -> b
$ String
"lookupQualNameEnv " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> QualName VName -> String
forall a. Show a => a -> String
show QualName VName
qn) (Maybe (TermScope dim) -> TermScope dim)
-> Maybe (TermScope dim) -> TermScope dim
forall a b. (a -> b) -> a -> b
$
    (StructType -> TypeBase dim NoMode)
-> TermScope dim -> QualName VName -> Maybe (TermScope dim)
forall dim.
(StructType -> TypeBase dim NoMode)
-> TermScope dim -> QualName VName -> Maybe (TermScope dim)
lookupQualNameEnvMaybe StructType -> TypeBase dim NoMode
onType TermScope dim
scope QualName VName
qn

-- | As 'lookupQualNameEnv', but 'Nothing' when the qualifiers do not name a
-- module in this scope. Use this when the name is merely being looked up
-- opportunistically, and so may well be one that has gone out of scope.
lookupQualNameEnvMaybe ::
  (StructType -> TypeBase dim NoMode) ->
  TermScope dim ->
  QualName VName ->
  Maybe (TermScope dim)
lookupQualNameEnvMaybe :: forall dim.
(StructType -> TypeBase dim NoMode)
-> TermScope dim -> QualName VName -> Maybe (TermScope dim)
lookupQualNameEnvMaybe StructType -> TypeBase dim NoMode
_ TermScope dim
scope (QualName [VName
q] VName
_)
  | VName -> Bool
isIntrinsic VName
q = TermScope dim -> Maybe (TermScope dim)
forall a. a -> Maybe a
Just TermScope dim
scope -- Magical intrinsic module.
lookupQualNameEnvMaybe StructType -> TypeBase dim NoMode
onType TermScope dim
scope (QualName [VName]
quals VName
_) = TermScope dim -> [VName] -> Maybe (TermScope dim)
descend TermScope dim
scope [VName]
quals
  where
    descend :: TermScope dim -> [VName] -> Maybe (TermScope dim)
descend TermScope dim
s [] = TermScope dim -> Maybe (TermScope dim)
forall a. a -> Maybe a
Just TermScope dim
s
    descend TermScope dim
s (VName
q : [VName]
qs)
      | Just (ModEnv Env
q_env) <- VName -> Map VName Mod -> Maybe Mod
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VName
q (Map VName Mod -> Maybe Mod) -> Map VName Mod -> Maybe Mod
forall a b. (a -> b) -> a -> b
$ TermScope dim -> Map VName Mod
forall dim. TermScope dim -> Map VName Mod
scopeModTable TermScope dim
s =
          TermScope dim -> [VName] -> Maybe (TermScope dim)
descend ((StructType -> TypeBase dim NoMode) -> Env -> TermScope dim
forall dim.
(StructType -> TypeBase dim NoMode) -> Env -> TermScope dim
envToTermScope StructType -> TypeBase dim NoMode
onType Env
q_env) [VName]
qs
      | Bool
otherwise =
          Maybe (TermScope dim)
forall a. Maybe a
Nothing

-- | An identifier corresponding to a type parameter, for size
-- parameters, which also exist as terms.
typeParamIdent :: TypeParam -> Maybe (Ident StructType)
typeParamIdent :: TypeParam -> Maybe (Ident StructType)
typeParamIdent (TypeParamDim VName
v SrcLoc
loc) =
  Ident StructType -> Maybe (Ident StructType)
forall a. a -> Maybe a
Just (Ident StructType -> Maybe (Ident StructType))
-> Ident StructType -> Maybe (Ident StructType)
forall a b. (a -> b) -> a -> b
$ VName -> Info StructType -> SrcLoc -> Ident StructType
forall {k} (f :: k -> *) vn (t :: k).
vn -> f t -> SrcLoc -> IdentBase f vn t
Ident VName
v (StructType -> Info StructType
forall a. a -> Info a
Info (StructType -> Info StructType) -> StructType -> Info StructType
forall a b. (a -> b) -> a -> b
$ ScalarTypeBase Size NoMode -> StructType
forall dim o. ScalarTypeBase dim o -> TypeBase dim o
Scalar (ScalarTypeBase Size NoMode -> StructType)
-> ScalarTypeBase Size NoMode -> StructType
forall a b. (a -> b) -> a -> b
$ PrimType -> ScalarTypeBase Size NoMode
forall dim o. PrimType -> ScalarTypeBase dim o
Prim (PrimType -> ScalarTypeBase Size NoMode)
-> PrimType -> ScalarTypeBase Size NoMode
forall a b. (a -> b) -> a -> b
$ IntType -> PrimType
Signed IntType
Int64) SrcLoc
loc
typeParamIdent TypeParam
_ = Maybe (Ident StructType)
forall a. Maybe a
Nothing