-- | Apply a substitution to a type
module Language.Giml.Types.Infer.Substitute where

import Control.Monad.Except
import Data.Generics.Uniplate.Data qualified as U
import Data.Map qualified as M
import Data.Set qualified as S
import Data.Text qualified as T
import Language.Giml.Syntax.Ast
import Language.Giml.Types.Infer.Types
import Language.Giml.Utils

--------------------------

-- * Substitute

-- | Apply a substitution.

-- ** API

-- | Replaces all type variables for any data type that
--   has an instance of Data using uniplate magic.
--
--   Note: uniplate magic = slow.
substitute :: (Substitute m) => (Data f) => Substitution -> f -> m f
substitute :: forall (m :: * -> *) f.
(Substitute m, Data f) =>
Substitution -> f -> m f
substitute Substitution
sub
  | Substitution -> Bool
forall k a. Map k a -> Bool
M.null Substitution
sub = f -> m f
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  | Bool
otherwise = (Type -> m Type) -> f -> m f
forall (m :: * -> *) from to.
(Monad m, Applicative m, Biplate from to) =>
(to -> m to) -> from -> m from
U.transformBiM (Substitution -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
Substitution -> Type -> m Type
replaceTypeVar Substitution
sub)

-- | Replaces all type variables for constraints.
substituteConstraints :: (Substitute m) => Substitution -> [ConstraintA] -> m [ConstraintA]
substituteConstraints :: forall (m :: * -> *).
Substitute m =>
Substitution -> [ConstraintA] -> m [ConstraintA]
substituteConstraints Substitution
sub
  | Substitution -> Bool
forall k a. Map k a -> Bool
M.null Substitution
sub = [ConstraintA] -> m [ConstraintA]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  | Bool
otherwise =
      (ConstraintA -> m ConstraintA) -> [ConstraintA] -> m [ConstraintA]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((ConstraintA -> m ConstraintA)
 -> [ConstraintA] -> m [ConstraintA])
-> (ConstraintA -> m ConstraintA)
-> [ConstraintA]
-> m [ConstraintA]
forall a b. (a -> b) -> a -> b
$ \case
        (Constraint
c, InputAnn
ann) ->
          let
            onConstraint :: (Type -> Type -> Constraint) -> Type -> Type -> m ConstraintA
onConstraint Type -> Type -> Constraint
constr Type
t1 Type
t2 =
              (Constraint -> ConstraintA) -> m Constraint -> m ConstraintA
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Constraint -> InputAnn -> ConstraintA)
-> InputAnn -> Constraint -> ConstraintA
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) InputAnn
ann) (m Constraint -> m ConstraintA) -> m Constraint -> m ConstraintA
forall a b. (a -> b) -> a -> b
$
                Type -> Type -> Constraint
constr
                  (Type -> Type -> Constraint) -> m Type -> m (Type -> Constraint)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Type -> m Type) -> Type -> m Type
forall (m :: * -> *). Monad m => (Type -> m Type) -> Type -> m Type
overTypes (Substitution -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
Substitution -> Type -> m Type
replaceTypeVar Substitution
sub) Type
t1
                  m (Type -> Constraint) -> m Type -> m Constraint
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Type -> m Type) -> Type -> m Type
forall (m :: * -> *). Monad m => (Type -> m Type) -> Type -> m Type
overTypes (Substitution -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
Substitution -> Type -> m Type
replaceTypeVar Substitution
sub) Type
t2
          in
            case Constraint
c of
              Equality Type
t1 Type
t2 ->
                (Type -> Type -> Constraint) -> Type -> Type -> m ConstraintA
onConstraint Type -> Type -> Constraint
Equality Type
t1 Type
t2

-- | Replaces all type variables for substitutions.
substituteSubs :: (Substitute m) => Substitution -> Substitution -> m Substitution
substituteSubs :: forall (m :: * -> *).
Substitute m =>
Substitution -> Substitution -> m Substitution
substituteSubs Substitution
sub
  | Substitution -> Bool
forall k a. Map k a -> Bool
M.null Substitution
sub = Substitution -> m Substitution
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  | Bool
otherwise = ((InputAnn, Type) -> m (InputAnn, Type))
-> Substitution -> m Substitution
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (((InputAnn, Type) -> m (InputAnn, Type))
 -> Substitution -> m Substitution)
-> ((InputAnn, Type) -> m (InputAnn, Type))
-> Substitution
-> m Substitution
forall a b. (a -> b) -> a -> b
$ (Type -> m Type) -> (InputAnn, Type) -> m (InputAnn, Type)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Type -> m Type) -> (InputAnn, Type) -> m (InputAnn, Type))
-> (Type -> m Type) -> (InputAnn, Type) -> m (InputAnn, Type)
forall a b. (a -> b) -> a -> b
$ (Type -> m Type) -> Type -> m Type
forall (m :: * -> *). Monad m => (Type -> m Type) -> Type -> m Type
overTypes (Substitution -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
Substitution -> Type -> m Type
replaceTypeVar Substitution
sub)

-- ** Types

-- | Monadic capabilities of Substitute
type Substitute m =
  ( MonadError TypeErrorA m
  )

-- ** Algorithm

-- | Find type variables that appear in the substitution and replace them.
replaceTypeVar :: (Substitute m) => Substitution -> Type -> m Type
replaceTypeVar :: forall (m :: * -> *).
Substitute m =>
Substitution -> Type -> m Type
replaceTypeVar Substitution
sub = \case
  TypeVar TypeVar
v ->
    m Type
-> ((InputAnn, Type) -> m Type) -> Maybe (InputAnn, Type) -> m Type
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
      (Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> Type
TypeVar TypeVar
v)
      ((InputAnn -> Type -> m Type) -> (InputAnn, Type) -> m Type
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((InputAnn -> Type -> m Type) -> (InputAnn, Type) -> m Type)
-> (InputAnn -> Type -> m Type) -> (InputAnn, Type) -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> InputAnn -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
TypeVar -> InputAnn -> Type -> m Type
occursCheck TypeVar
v) -- Make sure we don't have an infinite type
      (TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
v Substitution
sub)
  TypeRecExt [(Label, Type)]
r1 TypeVar
ext
    | Just (InputAnn
ann, Type
r2) <- TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
ext Substitution
sub -> do
        m Type -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Type -> m ()) -> m Type -> m ()
forall a b. (a -> b) -> a -> b
$ TypeVar -> InputAnn -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
TypeVar -> InputAnn -> Type -> m Type
occursCheck TypeVar
ext InputAnn
ann Type
r2
        InputAnn -> [(Label, Type)] -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
InputAnn -> [(Label, Type)] -> Type -> m Type
mergeRecords InputAnn
ann [(Label, Type)]
r1 Type
r2

  -- if the hidden type variable name is change we change it here as well
  TypePolyVariantUB TypeVar
tv [(Constr, Type)]
vars
    | Just (InputAnn
_, TypeVar TypeVar
tv') <- TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Substitution
sub -> do
        Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> [(Constr, Type)] -> Type
TypePolyVariantUB TypeVar
tv' [(Constr, Type)]
vars
  TypePolyVariantUB TypeVar
tv [(Constr, Type)]
vars1
    | Just (InputAnn
ann, Type
vars2) <- TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Substitution
sub -> do
        -- @TODO occursCheck ?
        InputAnn -> [(Constr, Type)] -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
InputAnn -> [(Constr, Type)] -> Type -> m Type
mergeVariants InputAnn
ann [(Constr, Type)]
vars1 Type
vars2

  -- this case represents a normal TypeVar
  TypePolyVariantLB [] TypeVar
tv
    | Just (InputAnn
_, Type
t) <- TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Substitution
sub -> do
        Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  TypePolyVariantLB [(Constr, Type)]
vars1 TypeVar
tv
    | Just (InputAnn
ann, Type
vars2) <- TypeVar -> Substitution -> Maybe (InputAnn, Type)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Substitution
sub -> do
        -- @TODO occursCheck ?
        InputAnn -> [(Constr, Type)] -> Type -> m Type
forall (m :: * -> *).
Substitute m =>
InputAnn -> [(Constr, Type)] -> Type -> m Type
mergeVariants InputAnn
ann [(Constr, Type)]
vars1 Type
vars2
  t :: Type
t@TypeCon {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypeApp {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypeRec {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypeRecExt {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypeVariant {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypePolyVariantLB {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypePolyVariantUB {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t
  t :: Type
t@TypeScheme {} -> Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t

-- Merging two open variants
mergeVariants :: (Substitute m) => InputAnn -> [(Constr, Type)] -> Type -> m Type
mergeVariants :: forall (m :: * -> *).
Substitute m =>
InputAnn -> [(Constr, Type)] -> Type -> m Type
mergeVariants InputAnn
ann [(Constr, Type)]
vars1 = \case
  TypePolyVariantLB [(Constr, Type)]
vars2 TypeVar
tv ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type)
-> (Map Constr Type -> Type) -> Map Constr Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(Constr, Type)] -> TypeVar -> Type)
-> TypeVar -> [(Constr, Type)] -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(Constr, Type)] -> TypeVar -> Type
TypePolyVariantLB TypeVar
tv ([(Constr, Type)] -> Type)
-> (Map Constr Type -> [(Constr, Type)]) -> Map Constr Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Constr Type -> [(Constr, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Constr Type -> m Type) -> Map Constr Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type -> Type) -> [(Constr, Type)] -> Map Constr Type
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Type -> Type -> Type) -> Type -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Type -> Type
forall a b. a -> b -> a
const) ([(Constr, Type)]
vars1 [(Constr, Type)] -> [(Constr, Type)] -> [(Constr, Type)]
forall a. Semigroup a => a -> a -> a
<> [(Constr, Type)]
vars2)
  TypePolyVariantUB TypeVar
tv [(Constr, Type)]
vars2 ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type)
-> (Map Constr Type -> Type) -> Map Constr Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeVar -> [(Constr, Type)] -> Type
TypePolyVariantUB TypeVar
tv ([(Constr, Type)] -> Type)
-> (Map Constr Type -> [(Constr, Type)]) -> Map Constr Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Constr Type -> [(Constr, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Constr Type -> m Type) -> Map Constr Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type -> Type) -> [(Constr, Type)] -> Map Constr Type
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Type -> Type -> Type) -> Type -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Type -> Type
forall a b. a -> b -> a
const) ([(Constr, Type)]
vars1 [(Constr, Type)] -> [(Constr, Type)] -> [(Constr, Type)]
forall a. Semigroup a => a -> a -> a
<> [(Constr, Type)]
vars2)
  TypeVariant [(Constr, Type)]
vars2 ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type)
-> (Map Constr Type -> Type) -> Map Constr Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Constr, Type)] -> Type
TypeVariant ([(Constr, Type)] -> Type)
-> (Map Constr Type -> [(Constr, Type)]) -> Map Constr Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Constr Type -> [(Constr, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Constr Type -> m Type) -> Map Constr Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type -> Type) -> [(Constr, Type)] -> Map Constr Type
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Type -> Type -> Type) -> Type -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Type -> Type
forall a b. a -> b -> a
const) ([(Constr, Type)]
vars1 [(Constr, Type)] -> [(Constr, Type)] -> [(Constr, Type)]
forall a. Semigroup a => a -> a -> a
<> [(Constr, Type)]
vars2)
  Type
t ->
    [InputAnn] -> TypeError -> m Type
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m Type) -> TypeError -> m Type
forall a b. (a -> b) -> a -> b
$ Type -> TypeError
NotAVariant Type
t

mergeRecords :: (Substitute m) => InputAnn -> [(Label, Type)] -> Type -> m Type
mergeRecords :: forall (m :: * -> *).
Substitute m =>
InputAnn -> [(Label, Type)] -> Type -> m Type
mergeRecords InputAnn
ann [(Label, Type)]
r1 = \case
  TypeRec [(Label, Type)]
r2 ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type)
-> (Map Label Type -> Type) -> Map Label Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Label, Type)] -> Type
TypeRec ([(Label, Type)] -> Type)
-> (Map Label Type -> [(Label, Type)]) -> Map Label Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Label Type -> [(Label, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Label Type -> m Type) -> Map Label Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type -> Type) -> [(Label, Type)] -> Map Label Type
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Type -> Type -> Type) -> Type -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Type -> Type
forall a b. a -> b -> a
const) ([(Label, Type)]
r1 [(Label, Type)] -> [(Label, Type)] -> [(Label, Type)]
forall a. Semigroup a => a -> a -> a
<> [(Label, Type)]
r2)
  TypeVar TypeVar
ext' ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ [(Label, Type)] -> TypeVar -> Type
TypeRecExt [(Label, Type)]
r1 TypeVar
ext'
  TypeRecExt [(Label, Type)]
r2 TypeVar
ext ->
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type)
-> (Map Label Type -> Type) -> Map Label Type -> m Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(Label, Type)] -> TypeVar -> Type)
-> TypeVar -> [(Label, Type)] -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip [(Label, Type)] -> TypeVar -> Type
TypeRecExt TypeVar
ext ([(Label, Type)] -> Type)
-> (Map Label Type -> [(Label, Type)]) -> Map Label Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Label Type -> [(Label, Type)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Label Type -> m Type) -> Map Label Type -> m Type
forall a b. (a -> b) -> a -> b
$ (Type -> Type -> Type) -> [(Label, Type)] -> Map Label Type
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith ((Type -> Type -> Type) -> Type -> Type -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip Type -> Type -> Type
forall a b. a -> b -> a
const) ([(Label, Type)]
r1 [(Label, Type)] -> [(Label, Type)] -> [(Label, Type)]
forall a. Semigroup a => a -> a -> a
<> [(Label, Type)]
r2)
  Type
t ->
    [InputAnn] -> TypeError -> m Type
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m Type) -> TypeError -> m Type
forall a b. (a -> b) -> a -> b
$ Type -> TypeError
NotARecord Type
t

-- | protect against infinite types
occursCheck :: (Substitute m) => TypeVar -> InputAnn -> Type -> m Type
occursCheck :: forall (m :: * -> *).
Substitute m =>
TypeVar -> InputAnn -> Type -> m Type
occursCheck TypeVar
v InputAnn
ann = \case
  TypeVar TypeVar
v'
    | TypeVar
v TypeVar -> TypeVar -> Bool
forall a. Eq a => a -> a -> Bool
== TypeVar
v' ->
        Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type -> m Type) -> Type -> m Type
forall a b. (a -> b) -> a -> b
$ TypeVar -> Type
TypeVar TypeVar
v'
  -- types that contain the type variable we are trying to replace
  -- are forbidden
  Type
t -> do
    let
      tvars :: [()]
tvars = [() | TypeVar TypeVar
tv <- Type -> [Type]
forall on. Uniplate on => on -> [on]
U.universe Type
t, TypeVar
tv TypeVar -> TypeVar -> Bool
forall a. Eq a => a -> a -> Bool
== TypeVar
v]
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([()] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [()]
tvars) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [InputAnn] -> TypeError -> m ()
forall (m :: * -> *) a.
MonadError TypeErrorA m =>
[InputAnn] -> TypeError -> m a
throwErr [InputAnn
ann] (TypeError -> m ()) -> TypeError -> m ()
forall a b. (a -> b) -> a -> b
$ TypeVar -> Type -> TypeError
InfiniteType TypeVar
v Type
t
    Type -> m Type
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
t

-- * Generalize

generalizeTermDef :: TermDef Ann -> TermDef Ann
generalizeTermDef :: TermDef Ann -> TermDef Ann
generalizeTermDef TermDef Ann
termdef =
  let
    t :: Type
t = Ann -> Type
annType (Ann -> Type) -> (TermDef Ann -> Ann) -> TermDef Ann -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TermDef Ann -> Ann
forall ann. TermDef ann -> ann
getTermAnn (TermDef Ann -> Type) -> TermDef Ann -> Type
forall a b. (a -> b) -> a -> b
$ TermDef Ann
termdef
    ([TypeVar]
freevars, Map TypeVar TypeVar
mapping) = Type -> ([TypeVar], Map TypeVar TypeVar)
generalizeType Type
t
    closed :: Type
closed = [TypeVar] -> Map TypeVar TypeVar -> Type -> Type
closeType [TypeVar]
freevars Map TypeVar TypeVar
mapping Type
t
    replace :: Expr Ann -> Expr Ann
replace = (TypeVar -> TypeVar) -> Expr Ann -> Expr Ann
forall from to. Biplate from to => (to -> to) -> from -> from
U.transformBi ((TypeVar -> TypeVar) -> Expr Ann -> Expr Ann)
-> (TypeVar -> TypeVar) -> Expr Ann -> Expr Ann
forall a b. (a -> b) -> a -> b
$ \TypeVar
tv ->
      TypeVar -> Maybe TypeVar -> TypeVar
forall a. a -> Maybe a -> a
fromMaybe TypeVar
tv (TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
mapping)
  in
    case TermDef Ann
termdef of
      Variable (Ann InputAnn
x Type
_) Var
name Maybe Type
_ Expr Ann
expr ->
        Ann -> Var -> Maybe Type -> Expr Ann -> TermDef Ann
forall a. a -> Var -> Maybe Type -> Expr a -> TermDef a
Variable (InputAnn -> Type -> Ann
Ann InputAnn
x Type
closed) Var
name (Type -> Maybe Type
forall a. a -> Maybe a
Just Type
closed) (Expr Ann -> Expr Ann
replace Expr Ann
expr)
      Function (Ann InputAnn
x Type
_) Var
name Maybe Type
_ [Maybe Var]
vars Expr Ann
body ->
        Ann -> Var -> Maybe Type -> [Maybe Var] -> Expr Ann -> TermDef Ann
forall a.
a -> Var -> Maybe Type -> [Maybe Var] -> Expr a -> TermDef a
Function (InputAnn -> Type -> Ann
Ann InputAnn
x Type
closed) Var
name (Type -> Maybe Type
forall a. a -> Maybe a
Just Type
closed) [Maybe Var]
vars (Expr Ann -> Expr Ann
replace Expr Ann
body)

generalizeAndClose :: Type -> Type
generalizeAndClose :: Type -> Type
generalizeAndClose Type
t = (([TypeVar] -> Map TypeVar TypeVar -> Type -> Type)
-> ([TypeVar], Map TypeVar TypeVar) -> Type -> Type
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [TypeVar] -> Map TypeVar TypeVar -> Type -> Type
closeType (([TypeVar], Map TypeVar TypeVar) -> Type -> Type)
-> ([TypeVar], Map TypeVar TypeVar) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ Type -> ([TypeVar], Map TypeVar TypeVar)
generalizeType Type
t) Type
t

generalizeType :: Type -> ([TypeVar], M.Map TypeVar TypeVar)
generalizeType :: Type -> ([TypeVar], Map TypeVar TypeVar)
generalizeType Type
t =
  let
    freevars :: [TypeVar]
freevars = Set TypeVar -> [TypeVar]
forall a. Set a -> [a]
S.toList (Set TypeVar -> [TypeVar]) -> Set TypeVar -> [TypeVar]
forall a b. (a -> b) -> a -> b
$ Type -> Set TypeVar
freeVars Type
t
    mapping :: Map TypeVar TypeVar
mapping =
      [(TypeVar, TypeVar)] -> Map TypeVar TypeVar
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(TypeVar, TypeVar)] -> Map TypeVar TypeVar)
-> [(TypeVar, TypeVar)] -> Map TypeVar TypeVar
forall a b. (a -> b) -> a -> b
$
        [TypeVar] -> [TypeVar] -> [(TypeVar, TypeVar)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TypeVar]
freevars ([TypeVar] -> [(TypeVar, TypeVar)])
-> [TypeVar] -> [(TypeVar, TypeVar)]
forall a b. (a -> b) -> a -> b
$
          (Char -> TypeVar) -> [Char] -> [TypeVar]
forall a b. (a -> b) -> [a] -> [b]
map (Var -> TypeVar
TV (Var -> TypeVar) -> (Char -> Var) -> Char -> TypeVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Var
T.pack ([Char] -> Var) -> (Char -> [Char]) -> Char -> Var
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> [Char]
forall (f :: * -> *) a. Applicative f => a -> f a
pure) [Char
'a' .. Char
'z'] [TypeVar] -> [TypeVar] -> [TypeVar]
forall a. Semigroup a => a -> a -> a
<> [Var -> TypeVar
TV (Var -> TypeVar) -> Var -> TypeVar
forall a b. (a -> b) -> a -> b
$ [Char] -> Var
T.pack (Char
'a' Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i) | Int
i :: Int <- [Int
0 ..]]
  in
    ([TypeVar]
freevars, Map TypeVar TypeVar
mapping)

closeType :: [TypeVar] -> M.Map TypeVar TypeVar -> Type -> Type
closeType :: [TypeVar] -> Map TypeVar TypeVar -> Type -> Type
closeType [TypeVar]
freevars Map TypeVar TypeVar
mapping Type
t =
  if [TypeVar] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TypeVar]
freevars
    then Type
t
    else [TypeVar] -> Type -> Type
TypeScheme (Map TypeVar TypeVar -> [TypeVar]
forall k a. Map k a -> [a]
M.elems Map TypeVar TypeVar
mapping) (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$
      ((TypeVar -> TypeVar) -> Type -> Type)
-> Type -> (TypeVar -> TypeVar) -> Type
forall a b c. (a -> b -> c) -> b -> a -> c
flip (TypeVar -> TypeVar) -> Type -> Type
forall from to. Biplate from to => (to -> to) -> from -> from
U.transformBi Type
t ((TypeVar -> TypeVar) -> Type) -> (TypeVar -> TypeVar) -> Type
forall a b. (a -> b) -> a -> b
$ \TypeVar
tv ->
        TypeVar -> Maybe TypeVar -> TypeVar
forall a. a -> Maybe a -> a
fromMaybe TypeVar
tv (TypeVar -> Map TypeVar TypeVar -> Maybe TypeVar
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup TypeVar
tv Map TypeVar TypeVar
mapping)

freeVars :: Type -> S.Set TypeVar
freeVars :: Type -> Set TypeVar
freeVars Type
t = [TypeVar] -> Set TypeVar
forall a. Ord a => [a] -> Set a
S.fromList [TypeVar
tv | TypeVar
tv <- Type -> [TypeVar]
forall from to. Biplate from to => from -> [to]
U.universeBi Type
t]