{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}
-- | Type utilities for Eidos: the total synthesizing 'typeOf', type
--   substitution (trivially capture-free: 'Ty' has no binders), arrow and
--   application spines, and type-level-natural evaluation.
--
--   'typeOf' follows the Core-Lint convention: it is total on well-formed
--   programs (doc/eidos.md §5) and calls 'error' on ill-formed ones, which
--   the linter rejects with a located diagnostic before any pass consults
--   types.
module ReWire.Eidos.Types
      ( typeOf
      , substTv
      , instantiate
      , mkArrow, dstArrow, flattenArrow
      , flattenTyApp, mkTyApp
      , flattenApp
      , evalNat, natNorm, tyEq
      , hasArrow, higherOrder, fundamental, reacOrStateT, synthable
      , machineDefn
      ) where

import ReWire.Annotation (Annote, ann)
import ReWire.Eidos.Syntax

import Data.HashMap.Strict (HashMap)
import Data.Text (Text)
import Numeric.Natural (Natural)

import qualified Data.HashMap.Strict as Map
import qualified Data.Text           as T

-- | Substitution of types for type variables. No renaming is ever needed:
--   types contain no binders.
substTv :: HashMap TyVar Ty -> Ty -> Ty
substTv :: HashMap TyVar Ty -> Ty -> Ty
substTv HashMap TyVar Ty
s = Ty -> Ty
go
      where go :: Ty -> Ty
            go :: Ty -> Ty
go = \ case
                  t :: Ty
t@(TyVarT Annote
_ TyVar
v)  -> Ty -> TyVar -> HashMap TyVar Ty -> Ty
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
Map.lookupDefault Ty
t TyVar
v HashMap TyVar Ty
s
                  TyApp Annote
an Ty
t Ty
u    -> Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Ty -> Ty
go Ty
t) (Ty -> Ty
go Ty
u)
                  Arrow Annote
an Ty
t Ty
u    -> Annote -> Ty -> Ty -> Ty
Arrow Annote
an (Ty -> Ty
go Ty
t) (Ty -> Ty
go Ty
u)
                  Ty
t               -> Ty
t

-- | Instantiate a signature at type arguments (must saturate the
--   quantifier list; the linter enforces this at occurrences).
instantiate :: Sig -> [Ty] -> Ty
instantiate :: Sig -> [Ty] -> Ty
instantiate (Sig [TyVar]
tvs Ty
t) [Ty]
ts
      | [TyVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVar]
tvs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Ty] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Ty]
ts = HashMap TyVar Ty -> Ty -> Ty
substTv ([(TyVar, Ty)] -> HashMap TyVar Ty
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
Map.fromList ([(TyVar, Ty)] -> HashMap TyVar Ty)
-> [(TyVar, Ty)] -> HashMap TyVar Ty
forall a b. (a -> b) -> a -> b
$ [TyVar] -> [Ty] -> [(TyVar, Ty)]
forall a b. [a] -> [b] -> [(a, b)]
zip [TyVar]
tvs [Ty]
ts) Ty
t
      | Bool
otherwise               = [Char] -> Ty
forall a. HasCallStack => [Char] -> a
error ([Char] -> Ty) -> [Char] -> Ty
forall a b. (a -> b) -> a -> b
$ [Char]
"Eidos.instantiate: signature quantifies "
            [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([TyVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVar]
tvs) [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" variables, applied to " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> [Char]
forall a. Show a => a -> [Char]
show ([Ty] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Ty]
ts)

-- | An arrow anywhere in the type.
hasArrow :: Ty -> Bool
hasArrow :: Ty -> Bool
hasArrow = \ case
      Arrow {}      -> Bool
True
      TyApp Annote
_ Ty
t1 Ty
t2 -> Ty -> Bool
hasArrow Ty
t1 Bool -> Bool -> Bool
|| Ty -> Bool
hasArrow Ty
t2
      Ty
_             -> Bool
False

-- | A function type with a function-typed parameter or result.
higherOrder :: Ty -> Bool
higherOrder :: Ty -> Bool
higherOrder (Ty -> ([Ty], Ty)
flattenArrow -> ([Ty]
doms, Ty
res)) = (Ty -> Bool) -> [Ty] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ty -> Bool
hasArrow ([Ty] -> Bool) -> [Ty] -> Bool
forall a b. (a -> b) -> a -> b
$ Ty
res Ty -> [Ty] -> [Ty]
forall a. a -> [a] -> [a]
: [Ty]
doms

-- | No String, Integer, or list constructors anywhere in the type.
fundamental :: Ty -> Bool
fundamental :: Ty -> Bool
fundamental = \ case
      TyCon Annote
_ TyConId
"String"  -> Bool
False
      TyCon Annote
_ TyConId
"Integer" -> Bool
False
      TyCon Annote
_ TyConId
"[_]"     -> Bool
False
      TyCon Annote
_ TyConId
"[]"      -> Bool
False
      Arrow Annote
_ Ty
t1 Ty
t2     -> Ty -> Bool
fundamental Ty
t1 Bool -> Bool -> Bool
&& Ty -> Bool
fundamental Ty
t2
      TyApp Annote
_ Ty
t1 Ty
t2     -> Ty -> Bool
fundamental Ty
t1 Bool -> Bool -> Bool
&& Ty -> Bool
fundamental Ty
t2
      Ty
_                 -> Bool
True

-- | Mentions a reactive-monad-stack constructor.
reacOrStateT :: Ty -> Bool
reacOrStateT :: Ty -> Bool
reacOrStateT = \ case
      TyCon Annote
_ TyConId
"ReacT"    -> Bool
True
      TyCon Annote
_ TyConId
"StateT"   -> Bool
True
      TyCon Annote
_ TyConId
"Identity" -> Bool
True
      Arrow Annote
_ Ty
t1 Ty
t2      -> Ty -> Bool
reacOrStateT Ty
t1 Bool -> Bool -> Bool
|| Ty -> Bool
reacOrStateT Ty
t2
      TyApp Annote
_ Ty
t1 Ty
t2      -> Ty -> Bool
reacOrStateT Ty
t1 Bool -> Bool -> Bool
|| Ty -> Bool
reacOrStateT Ty
t2
      Ty
_                  -> Bool
False

-- | Representable in hardware: first-order, fundamental, and no reactive
--   types in parameter position (the partial evaluator's per-definition
--   goal).
synthable :: Ty -> Bool
synthable :: Ty -> Bool
synthable Ty
t = Bool -> Bool
not (Ty -> Bool
higherOrder Ty
t)
           Bool -> Bool -> Bool
&& Ty -> Bool
fundamental Ty
t
           Bool -> Bool -> Bool
&& Bool -> Bool
not ((Ty -> Bool) -> [Ty] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Ty -> Bool
reacOrStateT ([Ty] -> Bool) -> [Ty] -> Bool
forall a b. (a -> b) -> a -> b
$ ([Ty], Ty) -> [Ty]
forall a b. (a, b) -> a
fst (([Ty], Ty) -> [Ty]) -> ([Ty], Ty) -> [Ty]
forall a b. (a -> b) -> a -> b
$ Ty -> ([Ty], Ty)
flattenArrow Ty
t)

mkArrow :: Annote -> Ty -> Ty -> Ty
mkArrow :: Annote -> Ty -> Ty -> Ty
mkArrow = Annote -> Ty -> Ty -> Ty
Arrow

dstArrow :: Ty -> Maybe (Ty, Ty)
dstArrow :: Ty -> Maybe (Ty, Ty)
dstArrow = \ case
      Arrow Annote
_ Ty
t Ty
u -> (Ty, Ty) -> Maybe (Ty, Ty)
forall a. a -> Maybe a
Just (Ty
t, Ty
u)
      Ty
_           -> Maybe (Ty, Ty)
forall a. Maybe a
Nothing

-- | An arrow spine: @flattenArrow (a -> b -> c) == ([a, b], c)@.
flattenArrow :: Ty -> ([Ty], Ty)
flattenArrow :: Ty -> ([Ty], Ty)
flattenArrow = \ case
      Arrow Annote
_ Ty
t Ty
u -> let ([Ty]
ts, Ty
r) = Ty -> ([Ty], Ty)
flattenArrow Ty
u in (Ty
t Ty -> [Ty] -> [Ty]
forall a. a -> [a] -> [a]
: [Ty]
ts, Ty
r)
      Ty
t           -> ([], Ty
t)

-- | A type application spine: @flattenTyApp (T a b) == (T, [a, b])@.
flattenTyApp :: Ty -> (Ty, [Ty])
flattenTyApp :: Ty -> (Ty, [Ty])
flattenTyApp = [Ty] -> Ty -> (Ty, [Ty])
go []
      where go :: [Ty] -> Ty -> (Ty, [Ty])
            go :: [Ty] -> Ty -> (Ty, [Ty])
go [Ty]
acc = \ case
                  TyApp Annote
_ Ty
t Ty
u -> [Ty] -> Ty -> (Ty, [Ty])
go (Ty
u Ty -> [Ty] -> [Ty]
forall a. a -> [a] -> [a]
: [Ty]
acc) Ty
t
                  Ty
t           -> (Ty
t, [Ty]
acc)

mkTyApp :: Annote -> Ty -> [Ty] -> Ty
mkTyApp :: Annote -> Ty -> [Ty] -> Ty
mkTyApp Annote
an = (Ty -> Ty -> Ty) -> Ty -> [Ty] -> Ty
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Ty -> Ty -> Ty) -> Ty -> [Ty] -> Ty)
-> (Ty -> Ty -> Ty) -> Ty -> [Ty] -> Ty
forall a b. (a -> b) -> a -> b
$ Annote -> Ty -> Ty -> Ty
TyApp Annote
an

-- | A term application spine: head and arguments, outermost last.
flattenApp :: Exp -> (Exp, [Arg])
flattenApp :: Exp -> (Exp, [Arg])
flattenApp = [Arg] -> Exp -> (Exp, [Arg])
go []
      where go :: [Arg] -> Exp -> (Exp, [Arg])
            go :: [Arg] -> Exp -> (Exp, [Arg])
go [Arg]
acc = \ case
                  App Annote
_ Exp
e Arg
a -> [Arg] -> Exp -> (Exp, [Arg])
go (Arg
a Arg -> [Arg] -> [Arg]
forall a. a -> [a] -> [a]
: [Arg]
acc) Exp
e
                  Exp
e         -> (Exp
e, [Arg]
acc)

-- | The type of an expression. Total on lint-clean programs; 'error' (with
--   the offending annotation) otherwise. Type arguments may only be applied
--   to 'Var' heads, must precede all term arguments, and must saturate the
--   head's quantifier list (doc/eidos.md §4.2). One tolerance: a bare
--   (argument-less) reference to a polymorphic name returns its open
--   signature type — the linter rejects such references outside
--   fully-instantiated spines, but the specializer reads them.
typeOf :: Exp -> Ty
typeOf :: Exp -> Ty
typeOf Exp
e = case Exp
e of
      Var Annote
_ Id
x         -> Sig -> [Ty] -> Ty
headTy (Id -> Sig
idSig Id
x) []
      Con Annote
_ Ty
t TyConId
_       -> Ty
t
      Prim Annote
_ Ty
t Builtin
_      -> Ty
t
      LitInt Annote
_ Ty
t Integer
_    -> Ty
t
      LitStr Annote
an TyConId
_     -> Annote -> TyConId -> Ty
TyCon Annote
an TyConId
"String"
      LitList Annote
_ Ty
t [Exp]
_   -> Ty
t
      LitVec Annote
_ Ty
t [Exp]
_    -> Ty
t
      Lam Annote
an Id
x Exp
b      -> Annote -> Ty -> Ty -> Ty
Arrow Annote
an (Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig Id
x) (Ty -> Ty) -> Ty -> Ty
forall a b. (a -> b) -> a -> b
$ Exp -> Ty
typeOf Exp
b
      Let Annote
_ Bind
_ Exp
b       -> Exp -> Ty
typeOf Exp
b
      Jump Annote
an JoinId
j [Exp]
args  -> Annote -> Int -> Ty -> Ty
peel Annote
an ([Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
args) (Ty -> Ty) -> Ty -> Ty
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ JoinId -> Id
jpId JoinId
j
      Case Annote
_ Ty
t Exp
_ Id
_ [Alt]
_  -> Ty
t
      App Annote
an Exp
_ Arg
_      -> Annote -> Ty
spineTy Annote
an
      where spineTy :: Annote -> Ty
            spineTy :: Annote -> Ty
spineTy Annote
an = case Exp -> (Exp, [Arg])
flattenApp Exp
e of
                  (Var Annote
_ Id
x, [Arg]
args)       -> Annote -> Ty -> [Exp] -> Ty
peelArgs Annote
an (Sig -> [Ty] -> Ty
headTy (Id -> Sig
idSig Id
x) ([Ty] -> Ty) -> [Ty] -> Ty
forall a b. (a -> b) -> a -> b
$ [Arg] -> [Ty]
targs [Arg]
args) ([Exp] -> Ty) -> [Exp] -> Ty
forall a b. (a -> b) -> a -> b
$ [Arg] -> [Exp]
eargs [Arg]
args
                  (Exp
h, [Arg]
args)
                        | [Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Arg] -> [Ty]
targs [Arg]
args) -> Annote -> Ty -> [Exp] -> Ty
peelArgs Annote
an (Exp -> Ty
typeOf Exp
h) ([Exp] -> Ty) -> [Exp] -> Ty
forall a b. (a -> b) -> a -> b
$ [Arg] -> [Exp]
eargs [Arg]
args
                        | Bool
otherwise         -> Annote -> [Char] -> Ty
forall a. Annote -> [Char] -> a
ill Annote
an [Char]
"type argument applied to a non-variable head"

            headTy :: Sig -> [Ty] -> Ty
            headTy :: Sig -> [Ty] -> Ty
headTy sig :: Sig
sig@(Sig [TyVar]
tvs Ty
t) [Ty]
ts
                  | [TyVar] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVar]
tvs Bool -> Bool -> Bool
&& [Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Ty]
ts       = Ty
t
                  | [TyVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVar]
tvs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Ty] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Ty]
ts   = Sig -> [Ty] -> Ty
instantiate Sig
sig [Ty]
ts
                  | [Ty] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Ty]
ts                   = Ty
t -- under-instantiated reference: free sig variables
                  | Bool
otherwise                 = Annote -> [Char] -> Ty
forall a. Annote -> [Char] -> a
ill (Exp -> Annote
forall a. Annotated a => a -> Annote
ann Exp
e) [Char]
"unsaturated type application"

            targs :: [Arg] -> [Ty]
            targs :: [Arg] -> [Ty]
targs [Arg]
as = [ Ty
t | TArg Ty
t <- (Arg -> Bool) -> [Arg] -> [Arg]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile Arg -> Bool
isTArg [Arg]
as ]

            eargs :: [Arg] -> [Exp]
            eargs :: [Arg] -> [Exp]
eargs [Arg]
as | (Arg -> Bool) -> [Arg] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Arg -> Bool
isTArg [Arg]
as' = Annote -> [Char] -> [Exp]
forall a. Annote -> [Char] -> a
ill (Exp -> Annote
forall a. Annotated a => a -> Annote
ann Exp
e) [Char]
"type argument after term arguments"
                     | Bool
otherwise      = [ Exp
x | EArg Exp
x <- [Arg]
as' ]
                  where as' :: [Arg]
as' = (Arg -> Bool) -> [Arg] -> [Arg]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Arg -> Bool
isTArg [Arg]
as

            isTArg :: Arg -> Bool
            isTArg :: Arg -> Bool
isTArg = \ case
                  TArg Ty
_ -> Bool
True
                  Arg
_      -> Bool
False

            peelArgs :: Annote -> Ty -> [Exp] -> Ty
            peelArgs :: Annote -> Ty -> [Exp] -> Ty
peelArgs Annote
an Ty
t [Exp]
es = Annote -> Int -> Ty -> Ty
peel Annote
an ([Exp] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exp]
es) Ty
t

            peel :: Annote -> Int -> Ty -> Ty
            peel :: Annote -> Int -> Ty -> Ty
peel Annote
_ Int
0 Ty
t  = Ty
t
            peel Annote
an Int
n Ty
t = case Ty -> Maybe (Ty, Ty)
dstArrow Ty
t of
                  Just (Ty
_, Ty
u) -> Annote -> Int -> Ty -> Ty
peel Annote
an (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Ty
u
                  Maybe (Ty, Ty)
Nothing     -> Annote -> [Char] -> Ty
forall a. Annote -> [Char] -> a
ill Annote
an [Char]
"term argument applied to a non-arrow"

            ill :: Annote -> String -> a
            ill :: forall a. Annote -> [Char] -> a
ill Annote
an [Char]
msg = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char] -> a) -> [Char] -> a
forall a b. (a -> b) -> a -> b
$ [Char]
"Eidos.typeOf: ill-formed expression (" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
msg [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
") at " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Annote -> [Char]
forall a. Show a => a -> [Char]
show Annote
an

-- | Evaluate a closed type-level natural: literals and the built-in
--   arithmetic constructors over evaluable operands.
evalNat :: Ty -> Maybe Natural
evalNat :: Ty -> Maybe Natural
evalNat Ty
t = case Ty
t of
      TyNat Annote
_ Natural
n -> Natural -> Maybe Natural
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Natural
n
      Ty
_         -> case Ty -> (Ty, [Ty])
flattenTyApp Ty
t of
            (TyCon Annote
_ TyConId
op, [Ty
a, Ty
b]) -> do
                  x <- Ty -> Maybe Natural
evalNat Ty
a
                  y <- evalNat b
                  natOp op x y
            (Ty, [Ty])
_                    -> Maybe Natural
forall a. Maybe a
Nothing
      where natOp :: Text -> Natural -> Natural -> Maybe Natural
            natOp :: TyConId -> Natural -> Natural -> Maybe Natural
natOp TyConId
op Natural
x Natural
y = case TyConId
op of
                  TyConId
"+" -> Natural -> Maybe Natural
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Maybe Natural) -> Natural -> Maybe Natural
forall a b. (a -> b) -> a -> b
$ Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
y
                  TyConId
"*" -> Natural -> Maybe Natural
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Maybe Natural) -> Natural -> Maybe Natural
forall a b. (a -> b) -> a -> b
$ Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
* Natural
y
                  TyConId
"-" | Natural
x Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
>= Natural
y -> Natural -> Maybe Natural
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural -> Maybe Natural) -> Natural -> Maybe Natural
forall a b. (a -> b) -> a -> b
$ Natural
x Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
y
                  TyConId
_   -> Maybe Natural
forall a. Maybe a
Nothing

-- | A definition the machine level keeps and the fold lowers: dotted-named
--   (not a builtin signature carrier, which is undotted), monomorphic, and
--   not reactive-typed (the reactive fragment is consumed by purification).
machineDefn :: Defn -> Bool
machineDefn :: Defn -> Bool
machineDefn Defn
d = (Char -> Bool) -> TyConId -> Bool
T.any (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'.') (Id -> TyConId
idOcc (Id -> TyConId) -> Id -> TyConId
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)
      Bool -> Bool -> Bool
&& [TyVar] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Sig -> [TyVar]
sigTVs (Sig -> [TyVar]) -> Sig -> [TyVar]
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)
      Bool -> Bool -> Bool
&& Bool -> Bool
not (Ty -> Bool
reacOrStateT (Ty -> Bool) -> Ty -> Bool
forall a b. (a -> b) -> a -> b
$ Sig -> Ty
sigTy (Sig -> Ty) -> Sig -> Ty
forall a b. (a -> b) -> a -> b
$ Id -> Sig
idSig (Id -> Sig) -> Id -> Sig
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d)

-- | Structural equality after normalization — the compiler-wide notion of
--   type equality (doc/eidos.md §5).
tyEq :: Ty -> Ty -> Bool
tyEq :: Ty -> Ty -> Bool
tyEq Ty
t Ty
t' = Ty -> Ty
natNorm Ty
t Ty -> Ty -> Bool
forall a. Eq a => a -> a -> Bool
== Ty -> Ty
natNorm Ty
t'

-- | Normalize a type: every type-level-natural subterm is put into a
--   canonical linear form (a sorted sum of coefficient-scaled atoms plus a
--   constant), so equal-modulo-arithmetic types — including OPEN ones,
--   e.g. @0 + m + n@ vs @n + m@, whose equality GHC's type-lits plugins
--   already proved — compare structurally equal. Type equality throughout
--   the compiler is structural equality after 'natNorm' (annotations are
--   already ignored by 'Eq Ty'). Subtraction does not distribute (naturals
--   truncate): @a - b@ normalizes its operands and is otherwise an atom.
natNorm :: Ty -> Ty
natNorm :: Ty -> Ty
natNorm Ty
t = case Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf Ty
t of
      Just (Natural, [(Ty, Natural)])
s  -> Annote -> (Natural, [(Ty, Natural)]) -> Ty
rebuild (Ty -> Annote
forall a. Annotated a => a -> Annote
ann Ty
t) (Natural, [(Ty, Natural)])
s
      Maybe (Natural, [(Ty, Natural)])
Nothing -> case Ty
t of
            TyApp Annote
an Ty
a Ty
b -> Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Ty -> Ty
natNorm Ty
a) (Ty -> Ty
natNorm Ty
b)
            Arrow Annote
an Ty
a Ty
b -> Annote -> Ty -> Ty -> Ty
Arrow Annote
an (Ty -> Ty
natNorm Ty
a) (Ty -> Ty
natNorm Ty
b)
            Ty
_            -> Ty
t

-- | A linear sum: constant + coefficient-scaled atoms (atoms canonically
--   ordered). Nothing when the type is not a natural-typed expression.
sumOf :: Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf :: Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf Ty
t = case Ty
t of
      TyNat Annote
_ Natural
n -> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural
n, [])
      Ty
_         -> case Ty -> (Ty, [Ty])
flattenTyApp Ty
t of
            (TyCon Annote
_ TyConId
"+", [Ty
a, Ty
b]) -> (Natural, [(Ty, Natural)])
-> (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)])
add ((Natural, [(Ty, Natural)])
 -> (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)]))
-> Maybe (Natural, [(Ty, Natural)])
-> Maybe ((Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)]))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf Ty
a Maybe ((Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)]))
-> Maybe (Natural, [(Ty, Natural)])
-> Maybe (Natural, [(Ty, Natural)])
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf Ty
b
            (TyCon Annote
_ TyConId
"*", [Ty
a, Ty
b]) -> do
                  sa <- Ty -> Maybe (Natural, [(Ty, Natural)])
sumOf Ty
a
                  sb <- sumOf b
                  mul sa sb
            (TyCon Annote
_ TyConId
"-", [Ty
a, Ty
b]) -> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural
0, [(Ty -> Ty -> Ty
subAtom Ty
a Ty
b, Natural
1)])
            (TyCon Annote
_ TyConId
op, [Ty]
_) | TyConId
op TyConId -> [TyConId] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ([TyConId
"+", TyConId
"*", TyConId
"-"] :: [Text]) -> Maybe (Natural, [(Ty, Natural)])
forall a. Maybe a
Nothing
            (Ty, [Ty])
_ | Ty -> Bool
isNatAtom Ty
t -> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural
0, [(Ty -> Ty
natNormAtom Ty
t, Natural
1)])
              | Bool
otherwise   -> Maybe (Natural, [(Ty, Natural)])
forall a. Maybe a
Nothing
      where add :: (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)])
            add :: (Natural, [(Ty, Natural)])
-> (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)])
add (Natural
c1, [(Ty, Natural)]
as1) (Natural
c2, [(Ty, Natural)]
as2) = (Natural
c1 Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
c2, [(Ty, Natural)] -> [(Ty, Natural)]
mergeAtoms ([(Ty, Natural)] -> [(Ty, Natural)])
-> [(Ty, Natural)] -> [(Ty, Natural)]
forall a b. (a -> b) -> a -> b
$ [(Ty, Natural)]
as1 [(Ty, Natural)] -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a. Semigroup a => a -> a -> a
<> [(Ty, Natural)]
as2)

            -- Only multiplication with at least one constant side stays
            -- linear; a product of two open sums is kept as an atom.
            mul :: (Natural, [(Ty, Natural)]) -> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
            mul :: (Natural, [(Ty, Natural)])
-> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
mul (Natural
c1, []) (Natural
c2, [(Ty, Natural)]
as2) = (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Natural
c1 Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
* Natural
c2, [ (Ty
a, Natural
c1 Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
* Natural
k) | (Ty
a, Natural
k) <- [(Ty, Natural)]
as2, Natural
c1 Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
0 ])
            mul (Natural
c1, [(Ty, Natural)]
as1) (Natural
c2, []) = (Natural, [(Ty, Natural)])
-> (Natural, [(Ty, Natural)]) -> Maybe (Natural, [(Ty, Natural)])
mul (Natural
c2, []) (Natural
c1, [(Ty, Natural)]
as1)
            mul (Natural, [(Ty, Natural)])
_ (Natural, [(Ty, Natural)])
_                = Maybe (Natural, [(Ty, Natural)])
forall a. Maybe a
Nothing

            subAtom :: Ty -> Ty -> Ty
            subAtom :: Ty -> Ty -> Ty
subAtom Ty
a Ty
b = Annote -> Ty -> Ty -> Ty
TyApp (Ty -> Annote
forall a. Annotated a => a -> Annote
ann Ty
t) (Annote -> Ty -> Ty -> Ty
TyApp (Ty -> Annote
forall a. Annotated a => a -> Annote
ann Ty
t) (Annote -> TyConId -> Ty
TyCon (Ty -> Annote
forall a. Annotated a => a -> Annote
ann Ty
t) TyConId
"-") (Ty -> Ty
natNorm Ty
a)) (Ty -> Ty
natNorm Ty
b)

            isNatAtom :: Ty -> Bool
            isNatAtom :: Ty -> Bool
isNatAtom = \ case
                  TyVarT Annote
_ TyVar
v  -> TyVar -> Kind
tvKind TyVar
v Kind -> Kind -> Bool
forall a. Eq a => a -> a -> Bool
== Kind
KNat
                  Ty
_           -> Bool
False

            natNormAtom :: Ty -> Ty
            natNormAtom :: Ty -> Ty
natNormAtom = Ty -> Ty
forall a. a -> a
id

mergeAtoms :: [(Ty, Natural)] -> [(Ty, Natural)]
mergeAtoms :: [(Ty, Natural)] -> [(Ty, Natural)]
mergeAtoms = ((Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)])
-> [(Ty, Natural)] -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
insert []
      where insert :: (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
            insert :: (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
insert (Ty
a, Natural
k) [] = [(Ty
a, Natural
k)]
            insert (Ty
a, Natural
k) ((Ty
b, Natural
j) : [(Ty, Natural)]
rest) = case Ty -> Ty -> Ordering
cmpTy Ty
a Ty
b of
                  Ordering
EQ -> (Ty
b, Natural
k Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
j) (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a. a -> [a] -> [a]
: [(Ty, Natural)]
rest
                  Ordering
LT -> (Ty
a, Natural
k) (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a. a -> [a] -> [a]
: (Ty
b, Natural
j) (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a. a -> [a] -> [a]
: [(Ty, Natural)]
rest
                  Ordering
GT -> (Ty
b, Natural
j) (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
forall a. a -> [a] -> [a]
: (Ty, Natural) -> [(Ty, Natural)] -> [(Ty, Natural)]
insert (Ty
a, Natural
k) [(Ty, Natural)]
rest

-- | Rebuild the canonical form: constant first (omitted when zero and atoms
--   exist), then coefficient-scaled atoms in canonical order, right-nested.
rebuild :: Annote -> (Natural, [(Ty, Natural)]) -> Ty
rebuild :: Annote -> (Natural, [(Ty, Natural)]) -> Ty
rebuild Annote
an (Natural
c, [(Ty, Natural)]
atoms) = case [Ty]
terms of
      []       -> Annote -> Natural -> Ty
TyNat Annote
an Natural
c
      (Ty
t : [Ty]
ts) -> (Ty -> Ty -> Ty) -> Ty -> [Ty] -> Ty
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\ Ty
acc Ty
u -> Ty -> Ty -> Ty
plus Ty
acc Ty
u) Ty
t [Ty]
ts
      where terms :: [Ty]
            terms :: [Ty]
terms = [ Annote -> Natural -> Ty
TyNat Annote
an Natural
c | Natural
c Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
0 Bool -> Bool -> Bool
|| [(Ty, Natural)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Ty, Natural)]
atoms ]
                  [Ty] -> [Ty] -> [Ty]
forall a. Semigroup a => a -> a -> a
<> [ Ty -> Natural -> Ty
scale Ty
a Natural
k | (Ty
a, Natural
k) <- [(Ty, Natural)]
atoms, Natural
k Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
0 ]

            plus :: Ty -> Ty -> Ty
            plus :: Ty -> Ty -> Ty
plus Ty
a Ty
b = Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Annote -> TyConId -> Ty
TyCon Annote
an TyConId
"+") Ty
a) Ty
b

            scale :: Ty -> Natural -> Ty
            scale :: Ty -> Natural -> Ty
scale Ty
a Natural
1 = Ty
a
            scale Ty
a Natural
k = Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Annote -> Ty -> Ty -> Ty
TyApp Annote
an (Annote -> TyConId -> Ty
TyCon Annote
an TyConId
"*") (Annote -> Natural -> Ty
TyNat Annote
an Natural
k)) Ty
a

-- | A deterministic structural ordering on types (annotations ignored;
--   variables by unique), for canonicalizing sums.
cmpTy :: Ty -> Ty -> Ordering
cmpTy :: Ty -> Ty -> Ordering
cmpTy Ty
a Ty
b = case (Ty
a, Ty
b) of
      (TyCon Annote
_ TyConId
c, TyCon Annote
_ TyConId
c')          -> TyConId -> TyConId -> Ordering
forall a. Ord a => a -> a -> Ordering
compare TyConId
c TyConId
c'
      (TyVarT Annote
_ TyVar
v, TyVarT Annote
_ TyVar
v')        -> TyVar -> TyVar -> Ordering
forall a. Ord a => a -> a -> Ordering
compare TyVar
v TyVar
v'
      (TyNat Annote
_ Natural
n, TyNat Annote
_ Natural
n')          -> Natural -> Natural -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Natural
n Natural
n'
      (TyApp Annote
_ Ty
t Ty
u, TyApp Annote
_ Ty
t' Ty
u')     -> Ty -> Ty -> Ordering
cmpTy Ty
t Ty
t' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Ty -> Ty -> Ordering
cmpTy Ty
u Ty
u'
      (Arrow Annote
_ Ty
t Ty
u, Arrow Annote
_ Ty
t' Ty
u')     -> Ty -> Ty -> Ordering
cmpTy Ty
t Ty
t' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Ty -> Ty -> Ordering
cmpTy Ty
u Ty
u'
      (Ty, Ty)
_                                -> Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Ty -> Int
tag Ty
a) (Ty -> Int
tag Ty
b)
      where tag :: Ty -> Int
            tag :: Ty -> Int
tag = \ case
                  TyCon {}  -> Int
0
                  TyApp {}  -> Int
1
                  TyVarT {} -> Int
2
                  TyNat {}  -> Int
3
                  Arrow {}  -> Int
4