{-# LANGUAGE LambdaCase #-}

{-|
Module      : Framework.Lambda.Types
Description : Curry typing with probabilistic types.
Copyright   : (c) Julian Grove and Aaron Steven White, 2025
License     : MIT
Maintainer  : julian.grove@gmail.com

Types and type inference are defined. Types feature variables, quantified at the
top level for some limited polymorphism. Constants may also be polymorphic.
-}

module Framework.Lambda.Types ( arity
                              , asTyped
                              , order
                              , Sig
                              , ty
                              , tyEq
                              , tyEq'
                              , Type(..)
                              , Typed(..)
                              , unify
                              ) where

import Control.Monad.State    (evalStateT, get, lift, put, StateT)
import Data.Char              (toLower)
import Data.List              (nub)
import Framework.Lambda.Terms

--------------------------------------------------------------------------------
-- * Types and type inference

-- ** Types

-- *** Definitions

-- | Arrows, products, and probabilistic types, as well as (a) abstract type
-- constructors, and (b) type variables for encoding polymorphism.
data Type = Atom String
          | Type :→ Type
          | Unit
          | Type :× Type
          | P Type
          | TyCon String [Type]
          | TyVar String
  deriving (Type -> Type -> Bool
(Type -> Type -> Bool) -> (Type -> Type -> Bool) -> Eq Type
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Type -> Type -> Bool
== :: Type -> Type -> Bool
$c/= :: Type -> Type -> Bool
/= :: Type -> Type -> Bool
Eq)

-- | Relations on types.
type TypeRel = Type -> Type -> Bool

-- | Type-level reductions.
type TypeRed = Type -> Maybe Type

-- | Determine if two types are semantically equivalent, taking type reductions
-- into account.
tyEq :: TypeRed -> TypeRel
tyEq :: TypeRed -> Type -> Type -> Bool
tyEq TypeRed
f Type
α Type
β = Constr -> Type -> Type
applySubst Constr
substα Type
α' Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Constr -> Type -> Type
applySubst Constr
substβ Type
β'
  where α', β' :: Type
        α' :: Type
α' = TypeRed -> Type -> Type
applyRed TypeRed
f Type
α
        β' :: Type
β' = TypeRed -> Type -> Type
applyRed TypeRed
f Type
β

        applyRed :: TypeRed -> Type -> Type
        applyRed :: TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty = case TypeRed
f Type
ty of
                         Just Type
ty' -> TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty'
                         Maybe Type
Nothing  -> TypeRed -> Type -> Type
descend TypeRed
f Type
ty

        descend :: TypeRed -> Type -> Type
        descend :: TypeRed -> Type -> Type
descend TypeRed
f = \case
          a :: Type
a@(Atom String
_)  -> Type
a
          Type
ty1 :→ Type
ty2  -> TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty1 Type -> Type -> Type
:→ TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty2
          Type
Unit        -> Type
Unit
          Type
ty1 :× Type
ty2  -> TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty1 Type -> Type -> Type
:× TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty2
          P Type
ty        -> Type -> Type
P (TypeRed -> Type -> Type
applyRed TypeRed
f Type
ty)
          TyCon String
g [Type]
tys -> String -> [Type] -> Type
TyCon String
g ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TypeRed -> Type -> Type
applyRed TypeRed
f) [Type]
tys)
          v :: Type
v@(TyVar String
_) -> Type
v

        substα, substβ :: Constr
        substα :: Constr
substα = [Type] -> [Type] -> Constr
forall a b. [a] -> [b] -> [(a, b)]
zip (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
αVars) (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
newVars)
        substβ :: Constr
substβ = [Type] -> [Type] -> Constr
forall a b. [a] -> [b] -> [(a, b)]
zip (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
βVars) (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
newVars)
  
        αVars, βVars, newVars :: TyVars
        αVars :: [String]
αVars = Type -> [String]
getVars Type
α'
        βVars :: [String]
βVars = Type -> [String]
getVars Type
β'
        newVars :: [String]
newVars = (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter ((String -> [String] -> Bool) -> [String] -> String -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
notElem ([String] -> String -> Bool) -> [String] -> String -> Bool
forall a b. (a -> b) -> a -> b
$ [String]
αVars [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String]
βVars) [String]
tyVars

-- | Match pop types to cancel out the types they pop.
popTy :: TypeRed
popTy :: TypeRed
popTy = \case
  TyCon (Char
'p' : Char
'o' : Char
'p' : String
f) [Type]
tys ->
    case [Type] -> Type
forall a. HasCallStack => [a] -> a
last [Type]
tys of
      TyCon String
g [Type]
tys'
        | String
f String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
g Bool -> Bool -> Bool
&&
          [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys') Bool -> Bool -> Bool
&&
          [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> Type -> Bool) -> [Type] -> [Type] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
(==) ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys) ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys'))  -> TypeRed
forall a. a -> Maybe a
Just ([Type] -> Type
forall a. HasCallStack => [a] -> a
last [Type]
tys')
      TyCon String
g [Type]
tys'
        | Bool
otherwise                                  ->
          (\Type
ty -> String -> [Type] -> Type
TyCon String
g ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys' [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type
ty])) (Type -> Type) -> Maybe Type -> Maybe Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
          TypeRed
popTy (String -> [Type] -> Type
TyCon (String
"pop" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
f) ([Type] -> [Type]
forall a. HasCallStack => [a] -> [a]
init [Type]
tys [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [[Type] -> Type
forall a. HasCallStack => [a] -> a
last [Type]
tys']))
      Type
_                                              -> Maybe Type
forall a. Maybe a
Nothing
  Type
_                               -> Maybe Type
forall a. Maybe a
Nothing

-- | Semantic type equality, taking popable types into account.
tyEq' :: TypeRel
tyEq' :: Type -> Type -> Bool
tyEq' = TypeRed -> Type -> Type -> Bool
tyEq TypeRed
popTy

instance Show Type where
  show :: Type -> String
show = \case
    Atom String
a                   -> String
a
    (a :: Type
a@(Type
_ :→ Type
_) :→ Type
b)        -> String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
") → " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b
    (Type
a :→ Type
b)                 -> Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" → " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b
    Type
Unit                     -> String
"⋄"
    a :: Type
a@(Type
_ :→ Type
_) :× b :: Type
b@(Type
_ :→ Type
_) -> String
"((" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
") × (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"))"
    a :: Type
a@(Type
_ :→ Type
_) :× Type
b          -> String
"((" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
") × " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
    Type
a :× b :: Type
b@(Type
_ :→ Type
_)          -> String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" × (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"))"
    Type
a :× Type
b                   -> Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" × " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
b
    P Type
a                      -> String
"P" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
wrap Type
a
    TyCon String
f [Type]
tys              -> String
f String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Type -> String) -> [Type] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Type -> String
wrap [Type]
tys
    TyVar String
s                  -> String
s
    where wrap :: Type -> String
          wrap :: Type -> String
wrap a :: Type
a@(Type
_ :→ Type
_)    = String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
          wrap a :: Type
a@(Type
_ :× Type
_)    = String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
          wrap a :: Type
a@(P Type
_)       = String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
          wrap a :: Type
a@(TyCon String
_ [Type]
_) = String
" (" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"
          wrap Type
a             = String
" "  String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
a

infixr 5 :→
infixl 6 :×

arity, order :: Type -> Int
arity :: Type -> Int
arity (Type
a :→ Type
b)      = Type -> Int
arity Type
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
arity Type
_             = Int
0
order :: Type -> Int
order (Type
a :→ Type
b)      = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Type -> Int
order Type
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Type -> Int
order Type
b)
order (Type
a :× Type
b)      = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Type -> Int
order Type
a) (Type -> Int
order Type
b)
order (P Type
a)         = Type -> Int
order Type
a
order (TyCon String
_ [Type]
tys) = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((Type -> Int) -> [Type] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Int
order [Type]
tys)
order Type
_             = Int
0
  
-- *** Type inference
  
-- | Lists of type equalities. We use these to find the most general unifier of
-- two principal types.
type Constr = [(Type, Type)]

-- | Type variable fs.
tyVars :: TyVars
tyVars :: [String]
tyVars = String
"" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (Integer -> String) -> [Integer] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> String
forall a. Show a => a -> String
show [Integer]
ints [String] -> (String -> [String]) -> [String]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
i -> (Char -> String) -> String -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Char -> String -> String
forall a. a -> [a] -> [a]
:String
i) [Char
'α'..Char
'ω']
  where ints :: [Integer]
        ints :: [Integer]
ints = Integer
1 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Integer
forall a. Enum a => a -> a
succ [Integer]
ints

-- | Get the variable fs that occur in some type.
getVars :: Type -> TyVars
getVars :: Type -> [String]
getVars = [String] -> [String]
forall a. Eq a => [a] -> [a]
nub ([String] -> [String]) -> (Type -> [String]) -> Type -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> [String]
getVars'
  where getVars' :: Type -> TyVars
        getVars' :: Type -> [String]
getVars' (TyVar String
v)     = [String
v]
        getVars' (Type
a :→ Type
b)      = Type -> [String]
getVars' Type
a [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ Type -> [String]
getVars' Type
b
        getVars' (Type
a :× Type
b)      = Type -> [String]
getVars' Type
a [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ Type -> [String]
getVars' Type
b
        getVars' (P Type
a)         = Type -> [String]
getVars' Type
a
        getVars' (TyCon String
_ [Type]
tys) = (Type -> [String]) -> [Type] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Type -> [String]
getVars' [Type]
tys
        getVars' Type
_             = []

-- | Arrive at a most general unifier, or fail if none exists.
unify :: Constr -> Maybe Constr
unify :: Constr -> Maybe Constr
unify Constr
cs = do
  (Constr
cs', []) <- (Constr, Constr) -> Maybe (Constr, Constr)
pass ([], Constr
cs)
  if Constr
cs' Constr -> Constr -> Bool
forall a. Eq a => a -> a -> Bool
== Constr
cs then Constr -> Maybe Constr
forall a. a -> Maybe a
Just Constr
cs else Constr -> Maybe Constr
unify Constr
cs'
  where
    -- | Make one pass through the whole list.
    pass :: (Constr, Constr) -> Maybe (Constr, Constr)
    pass :: (Constr, Constr) -> Maybe (Constr, Constr)
pass = \case
      p :: (Constr, Constr)
p@(Constr
_, [])  -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr, Constr)
p
      (Constr
u1, (Type, Type)
x:Constr
u2) -> do
        (Constr
u1', Constr
u2') <- (Constr, Constr) -> (Type, Type) -> Maybe (Constr, Constr)
step (Constr
u1, Constr
u2) (Type, Type)
x
        (Constr, Constr) -> Maybe (Constr, Constr)
pass (Constr
u1', Constr
u2')

    -- | Take one step toward a most general unifier.
    step :: (Constr, Constr) -> (Type, Type) -> Maybe (Constr, Constr)
    step :: (Constr, Constr) -> (Type, Type) -> Maybe (Constr, Constr)
step (Constr
u1, Constr
u2) = \case
      (Type
x, Type
y)              | Type
x Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
y     -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1, Constr
u2) -- Get rid of trivial equalities.
      (x :: Type
x@(TyVar String
_), Type
y)    | Type -> Type -> Bool
occurs Type
x Type
y -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- Impredicativity; uh oh!
      eq :: (Type, Type)
eq@(x :: Type
x@(TyVar String
_), Type
y)              -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just ( ((Type, Type) -> (Type, Type)) -> Constr -> Constr
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> (Type, Type) -> (Type, Type)
substEqs Type
x Type
y) Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type, Type)
eq]
                                               , ((Type, Type) -> (Type, Type)) -> Constr -> Constr
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> (Type, Type) -> (Type, Type)
substEqs Type
x Type
y) Constr
u2 ) -- Substitute.
      (Type
x, TyVar String
y)                     -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(String -> Type
TyVar String
y, Type
x)], Constr
u2) -- Swap.
      (Atom String
_, Type
_)                      -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- We've already handled equalities.
      (Type
x :→ Type
y, Type
z :→ Type
w)                 -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type
x, Type
z), (Type
y, Type
w)], Constr
u2) -- Break apart arrows.
      (Type
_ :→ Type
_, Type
_)                      -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- No mismatches allowed.
      (Type
Unit, Type
_)                        -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- No mismatches allowed.
      (Type
x :× Type
y, Type
z :× Type
w)                 -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type
x, Type
z), (Type
y, Type
w)], Constr
u2) -- Break apart products
      (Type
_ :× Type
_, Type
_)                      -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- No mismatches allowed.
      (P Type
x, P Type
y)                       -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type
x, Type
y)], Constr
u2) -- Break apart @P@.
      (P Type
_, Type
_)                         -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- No mismatches allowed.
      (TyCon String
f [Type]
tys1
        , TyCon String
g [Type]
tys2)  | String
f String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
g Bool -> Bool -> Bool
&&
                           [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==
                           [Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
tys2 -> (Constr, Constr) -> Maybe (Constr, Constr)
forall a. a -> Maybe a
Just (Constr
u1 Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [Type] -> [Type] -> Constr
forall a b. [a] -> [b] -> [(a, b)]
zip [Type]
tys1 [Type]
tys2, Constr
u2) -- Break apart @TyCon@.
      (TyCon String
_ [Type]
_, Type
_)                   -> Maybe (Constr, Constr)
forall a. Maybe a
Nothing -- No mismatches allowed.

    -- | Substitute inside equalities.
    substEqs :: Type -> Type -> (Type, Type) -> (Type, Type)
    substEqs :: Type -> Type -> (Type, Type) -> (Type, Type)
substEqs Type
x Type
y = \(Type
z, Type
w) -> (Type -> Type -> Type -> Type
substType Type
x Type
y Type
z, Type -> Type -> Type -> Type
substType Type
x Type
y Type
w)

    -- | Substitute inside a type.
    substType :: Type -> Type -> Type -> Type
    substType :: Type -> Type -> Type -> Type
substType Type
x Type
y Type
z             | Type
z Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
x = Type
y
    substType Type
x Type
y (Type
z1 :→ Type
z2)             = Type -> Type -> Type -> Type
substType Type
x Type
y Type
z1 Type -> Type -> Type
:→ Type -> Type -> Type -> Type
substType Type
x Type
y Type
z2
    substType Type
x Type
y (Type
z1 :× Type
z2)             = Type -> Type -> Type -> Type
substType Type
x Type
y Type
z1 Type -> Type -> Type
:× Type -> Type -> Type -> Type
substType Type
x Type
y Type
z2
    substType Type
x Type
y (P Type
z)                  = Type -> Type
P (Type -> Type -> Type -> Type
substType Type
x Type
y Type
z)
    substType Type
x Type
y (TyCon String
f [Type]
tys)          = String -> [Type] -> Type
TyCon String
f ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> Type -> Type
substType Type
x Type
y) [Type]
tys)
    substType Type
_ Type
_ Type
z                      = Type
z

    -- | Does @x@ occur in @y@?
    occurs :: Type -> Type -> Bool
    occurs :: Type -> Type -> Bool
occurs Type
x = \case -- List out possible @y@s.
      Type
y :→ Type
z      -> Type -> Type -> Bool
occurs Type
x Type
y Bool -> Bool -> Bool
|| Type -> Type -> Bool
occurs Type
x Type
z
      Type
y :× Type
z      -> Type -> Type -> Bool
occurs Type
x Type
y Bool -> Bool -> Bool
|| Type -> Type -> Bool
occurs Type
x Type
z
      P Type
y         -> Type -> Type -> Bool
occurs Type
x Type
y
      TyCon String
_ [Type]
tys -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> Type -> Bool
occurs Type
x) [Type]
tys)
      Type
y           -> Type
x Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
y -- The "otherwise" case; fly, @At@, @Unit@, and @TyVar@.

-- | Assign types to constants.
type Sig = Constant -> Maybe Type

type TyVars    = [String]
type VarTyping = String -> Maybe Type
type TVC       = (TyVars, VarTyping, Constr)

-- | Compute the principal type of a term, given a signature.
computeType :: Sig -> Term -> StateT TVC Maybe Type
computeType :: Sig -> Term -> StateT TVC Maybe Type
computeType Sig
tau Term
t = do Type
ty <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
                       -- | Get type constraints.
                       ([String]
_, VarTyping
_, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
                       -- | Find a most general unifier.
                       Constr
subst <- Maybe Constr -> StateT TVC Maybe Constr
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Constr -> Maybe Constr
unify Constr
cs)
                       -- | Use it to substitute.
                       Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Constr -> Type -> Type
applySubst Constr
subst Type
ty)
  where collectConstraints :: Sig -> Term -> StateT TVC Maybe Type
        collectConstraints :: Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau = \case
          Var String
v -> do
            (String
ty:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            -- | Check if @v@ is already assigned a type by @f@. If not, assign
            -- it a fresh type variable; if so, use the type it is assigned.
            let g :: String -> Maybe Type
                g :: VarTyping
g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v = case Maybe Type
fv of
                                   Maybe Type
Nothing    -> TypeRed
forall a. a -> Maybe a
Just (String -> Type
TyVar String
ty)
                                   x :: Maybe Type
x@(Just Type
_) -> Maybe Type
x
                g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
v = VarTyping
f String
v'

                types :: TyVars
                types :: [String]
types = case Maybe Type
fv of
                          Just Type
_  -> String
tyString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
pes
                          Maybe Type
Nothing -> [String]
pes
                
                fv :: Maybe Type
                fv :: Maybe Type
fv = VarTyping
f String
v
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
types, VarTyping
g, Constr
cs)
            Maybe Type -> StateT TVC Maybe Type
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (VarTyping
g String
v) -- Return the (possibly new) type assigned to @v@.
          Con Constant
c -> do
            ([String]
types, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            Type
cType <- Maybe Type -> StateT TVC Maybe Type
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Sig
tau Constant
c)
            -- | Ensure that the type variables used for polymorphic constants
            -- are fresh (to prevent any potential clashes), and unify them with
            -- the original type variables.
            let newVars :: [String]
newVars = Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
take Int
lng [String]
types'
                rest :: [String]
rest    = Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
drop Int
lng [String]
types'
                types' :: [String]
types'  = (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter (String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Type -> [String]
getVars Type
cType) [String]
types
                lng :: Int
lng     = [String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Type -> [String]
getVars Type
cType)
                conCs :: Constr
conCs   = [Type] -> [Type] -> Constr
forall a b. [a] -> [b] -> [(a, b)]
zip (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Type -> [String]
getVars Type
cType) (String -> Type
TyVar (String -> Type) -> [String] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
newVars)
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
rest, VarTyping
f, Constr
cs)
            Constr
subst <- Maybe Constr -> StateT TVC Maybe Constr
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Constr -> Maybe Constr
unify Constr
conCs) 
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Constr -> Type -> Type
applySubst Constr
subst Type
cType)
          App Term
t Term
u -> do
            -- | Process the argument type; process the function type; then
            -- constrain them to be consistent.
            Type
uType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
u
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            (String
ty:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes, VarTyping
f, Constr
cs Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type
uType Type -> Type -> Type
:→ String -> Type
TyVar String
ty, Type
tType)])
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Type
TyVar String
ty)
          Lam String
v Term
t -> do
            -- | Get rid of any type assignments to the abstracted variable
            -- while the body of the abstraction is processed.
            (String
ty:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            let g :: VarTyping
                g :: VarTyping
g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v = TypeRed
forall a. a -> Maybe a
Just (String -> Type
TyVar String
ty)
                g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
v = VarTyping
f String
v'
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes, VarTyping
g, Constr
cs)
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            ([String]
types, VarTyping
h, Constr
cs') <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            Type
newTy <- Maybe Type -> StateT TVC Maybe Type
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (VarTyping
g String
v)
            -- | Revert the outgoing variable-typing scheme to the original one
            -- after typing the abstration.
            let i :: VarTyping
                i :: VarTyping
i String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v = VarTyping
f String
v'
                i String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
v = VarTyping
h String
v'
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
types, VarTyping
i, Constr
cs')
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
newTy Type -> Type -> Type
:→ Type
tType)
          Term
TT -> Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
Unit
          Pi1 Term
t -> do
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            (String
t:String
y:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes, VarTyping
f , Constr
cs Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(String -> Type
TyVar String
t Type -> Type -> Type
:× String -> Type
TyVar String
y, Type
tType)])
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Type
TyVar String
t)
          Pi2 Term
t -> do
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            (String
t:String
y:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes, VarTyping
f, Constr
cs Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(String -> Type
TyVar String
t Type -> Type -> Type
:× String -> Type
TyVar String
y, Type
tType)])
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Type
TyVar String
y)
          Pair Term
t Term
u -> do
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            Type
uType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
u
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
tType Type -> Type -> Type
:× Type
uType)
          Let String
v Term
t Term
u -> do
            -- | Get rid of any type assignments to the let-bound variable while
            -- its scope is processed.
            (String
ty:[String]
pes, VarTyping
f, Constr
cs) <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            let g :: VarTyping
                g :: VarTyping
g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v = TypeRed
forall a. a -> Maybe a
Just (String -> Type
TyVar String
ty)
                g String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
v = VarTyping
f String
v'
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes, VarTyping
g, Constr
cs)
            Type
uType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
u
            ([String]
types, VarTyping
h, Constr
cs') <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            Type
newTy <- Maybe Type -> StateT TVC Maybe Type
forall (m :: * -> *) a. Monad m => m a -> StateT TVC m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (VarTyping
h String
v)
            -- | Revert the outgoing variable-typing scheme to the original one
            -- after typing the let-construct.
            let i :: VarTyping
                i :: VarTyping
i String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v = VarTyping
f String
v'
                i String
v' | String
v' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
/= String
v = VarTyping
h String
v'
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
types, VarTyping
i, Constr
cs')
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            (String
ty':[String]
pes', VarTyping
j, Constr
cs'') <- StateT TVC Maybe TVC
forall s (m :: * -> *). MonadState s m => m s
get
            TVC -> StateT TVC Maybe ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([String]
pes', VarTyping
j, Constr
cs'' Constr -> Constr -> Constr
forall a. [a] -> [a] -> [a]
++ [(Type
tType, Type -> Type
P Type
newTy), (Type
uType, Type -> Type
P (String -> Type
TyVar String
ty'))])
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
uType
          Return Term
t -> do
            Type
tType <- Sig -> Term -> StateT TVC Maybe Type
collectConstraints Sig
tau Term
t
            Type -> StateT TVC Maybe Type
forall a. a -> StateT TVC Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> Type
P Type
tType)

-- | Apply a substitution to a type, given a set of constraints.
applySubst :: Constr -> Type -> Type
applySubst :: Constr -> Type -> Type
applySubst Constr
s = \case
  a :: Type
a@(Atom String
_)  -> Type
a
  Type
x :→ Type
y      -> Constr -> Type -> Type
applySubst Constr
s Type
x Type -> Type -> Type
:→ Constr -> Type -> Type
applySubst Constr
s Type
y
  Type
Unit        -> Type
Unit
  Type
x :× Type
y      -> Constr -> Type -> Type
applySubst Constr
s Type
x Type -> Type -> Type
:× Constr -> Type -> Type
applySubst Constr
s Type
y
  P Type
x         -> Type -> Type
P (Constr -> Type -> Type
applySubst Constr
s Type
x)
  TyCon String
f [Type]
tys -> String -> [Type] -> Type
TyCon String
f ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Constr -> Type -> Type
applySubst Constr
s) [Type]
tys)
  TyVar String
v     -> String -> Constr -> Type
lookUp String
v Constr
s
    where  lookUp :: String -> Constr -> Type
           lookUp :: String -> Constr -> Type
lookUp String
v  []                   = String -> Type
TyVar String
v
           lookUp String
v1 ((TyVar String
v2, Type
t) : Constr
s') = if   String
v2 String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
v1
                                            then Type
t
                                            else String -> Constr -> Type
lookUp String
v1 Constr
s'

-- | Typed terms (where not every term has a type).
data Typed = Typed { Typed -> Term
termOf :: Term, Typed -> Maybe Type
typeOf :: Maybe Type } deriving (Typed -> Typed -> Bool
(Typed -> Typed -> Bool) -> (Typed -> Typed -> Bool) -> Eq Typed
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Typed -> Typed -> Bool
== :: Typed -> Typed -> Bool
$c/= :: Typed -> Typed -> Bool
/= :: Typed -> Typed -> Bool
Eq)

instance Show Typed where
  show :: Typed -> String
show (Typed Term
te (Just Type
ty)) = Term -> String
forall a. Show a => a -> String
show Term
te String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" : " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
ty
  show (Typed Term
te Maybe Type
Nothing)   = String
"⌜" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Term -> String
forall a. Show a => a -> String
show Term
te String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"⌝ is not assigned a type."

-- | Provide a term along with its principal type, given a signature.
ty :: Sig -> Term -> Typed
ty :: Sig -> Term -> Typed
ty Sig
tau Term
te = Term -> Maybe Type -> Typed
Typed Term
te Maybe Type
theType
    where theType :: Maybe Type
          theType :: Maybe Type
theType = StateT TVC Maybe Type -> TVC -> Maybe Type
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT (Sig -> Term -> StateT TVC Maybe Type
computeType Sig
tau Term
te) ([String]
tyVars, VarTyping
start, Constr
empty)

          tyVars :: TyVars
          tyVars :: [String]
tyVars = String
"" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (Integer -> String) -> [Integer] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> String
forall a. Show a => a -> String
show [Integer]
ints [String] -> (String -> [String]) -> [String]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \String
i -> (Char -> String) -> String -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Char -> String -> String
forall a. a -> [a] -> [a]
:String
i) [Char
'α'..Char
'ω']

          ints :: [Integer]
          ints :: [Integer]
ints = Integer
1 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Integer
forall a. Enum a => a -> a
succ [Integer]
ints

          start :: VarTyping
          start :: VarTyping
start = Maybe Type -> VarTyping
forall a b. a -> b -> a
const Maybe Type
forall a. Maybe a
Nothing

          empty :: Constr
          empty :: Constr
empty = []

asTyped :: Sig -> (Term -> Term) -> Typed -> Typed
asTyped :: Sig -> (Term -> Term) -> Typed -> Typed
asTyped Sig
tau Term -> Term
f = Sig -> Term -> Typed
ty Sig
tau (Term -> Typed) -> (Typed -> Term) -> Typed -> Typed
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Term
f (Term -> Term) -> (Typed -> Term) -> Typed -> Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Typed -> Term
termOf