content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TupleSections     #-}
  
  -- | Constraint-based typing from the presentation in Pierce's book.
  module Kempe.TyAssign ( TypeM
                        , runTypeM
                        , checkModule
                        , assignModule
                        ) where
  
  import           Control.Composition        (thread, (.$))
  import           Control.Monad              (foldM, replicateM, unless, when, zipWithM_)
  import           Control.Monad.Except       (throwError)
  import           Control.Monad.State.Strict (State, StateT, evalState, get, gets, modify, put, runStateT)
  import           Data.Bifunctor             (bimap, second)
  import           Data.Foldable              (traverse_)
  import           Data.Functor               (void, ($>))
  import qualified Data.IntMap                as IM
  import           Data.List.NonEmpty         (NonEmpty (..))
  import qualified Data.Set                   as S
  import qualified Data.Text                  as T
  import           Data.Tuple.Ext             (fst3)
  import           Kempe.AST
  import           Kempe.Error
  import           Kempe.Name
  import           Kempe.Unique
  import           Lens.Micro                 (Lens', over)
  import           Lens.Micro.Mtl             (modifying, (.=))
  import           Prettyprinter              (Doc, Pretty (pretty), vsep, (<+>))
  import           Prettyprinter.Debug
  import           Prettyprinter.Ext
  
  type TyEnv a = IM.IntMap (StackType a)
  
  data TyState a = TyState { maxU             :: Int -- ^ For renamer
                           , tyEnv            :: TyEnv a
                           , kindEnv          :: IM.IntMap Kind
                           , renames          :: IM.IntMap Int
                           , constructorTypes :: IM.IntMap (StackType a)
                           , constraints      :: S.Set (KempeTy a, KempeTy a) -- Just need equality between simple types? (do have tyapp but yeah)
                           }
  
  instance Pretty (TyState a) where
      pretty (TyState _ te _ r _ cs) =
          "type environment:" <#> vsep (prettyBound <$> IM.toList te)
              <#> "renames:" <#*> prettyDumpBinds r
              <#> "constraints:" <#> prettyConstraints cs
  
  prettyConstraints :: S.Set (KempeTy a, KempeTy a) -> Doc ann
  prettyConstraints cs = vsep (prettyEq <$> S.toList cs)
  
  prettyEq :: (KempeTy a, KempeTy a) -> Doc ann
  prettyEq (ty, ty') = pretty ty <+> "≡" <+> pretty ty'
  
  prettyDumpBinds :: Pretty b => IM.IntMap b -> Doc a
  prettyDumpBinds b = vsep (prettyBind <$> IM.toList b)
  
  emptyStackType :: StackType a
  emptyStackType = StackType [] []
  
  maxULens :: Lens' (TyState a) Int
  maxULens f s = fmap (\x -> s { maxU = x }) (f (maxU s))
  
  constructorTypesLens :: Lens' (TyState a) (IM.IntMap (StackType a))
  constructorTypesLens f s = fmap (\x -> s { constructorTypes = x }) (f (constructorTypes s))
  
  tyEnvLens :: Lens' (TyState a) (TyEnv a)
  tyEnvLens f s = fmap (\x -> s { tyEnv = x }) (f (tyEnv s))
  
  kindEnvLens :: Lens' (TyState a) (IM.IntMap Kind)
  kindEnvLens f s = fmap (\x -> s { kindEnv = x }) (f (kindEnv s))
  
  renamesLens :: Lens' (TyState a) (IM.IntMap Int)
  renamesLens f s = fmap (\x -> s { renames = x }) (f (renames s))
  
  constraintsLens :: Lens' (TyState a) (S.Set (KempeTy a, KempeTy a))
  constraintsLens f s = fmap (\x -> s { constraints = x }) (f (constraints s))
  
  dummyName :: T.Text -> TypeM () (Name ())
  dummyName n = do
      pSt <- gets maxU
      Name n (Unique $ pSt + 1) ()
          <$ modifying maxULens (+1)
  
  data Kind = Star
            | TyCons Kind Kind
            deriving (Eq)
  
  type TypeM a = StateT (TyState a) (Either (Error a))
  
  type UnifyMap = IM.IntMap (KempeTy ())
  
  inContext :: UnifyMap -> KempeTy () -> KempeTy ()
  inContext um ty'@(TyVar _ (Name _ (Unique i) _)) =
      case IM.lookup i um of
          Just ty@TyVar{} -> inContext (IM.delete i um) ty -- prevent cyclic lookups
          -- TODO: does this need a case for TyApp -> inContext?
          Just ty         -> ty
          Nothing         -> ty'
  inContext _ ty'@TyBuiltin{} = ty'
  inContext _ ty'@TyNamed{} = ty'
  inContext um (TyApp l ty ty') = TyApp l (inContext um ty) (inContext um ty')
  
  -- | Perform substitutions before handing off to 'unifyMatch'
  unifyPrep :: UnifyMap
            -> [(KempeTy (), KempeTy ())]
            -> Either (Error ()) (IM.IntMap (KempeTy ()))
  unifyPrep _ [] = Right mempty
  unifyPrep um ((ty, ty'):tys) =
      let ty'' = inContext um ty
          ty''' = inContext um ty'
      in unifyMatch um $ (ty'', ty'''):tys
  
  unifyMatch :: UnifyMap -> [(KempeTy (), KempeTy ())] -> Either (Error ()) (IM.IntMap (KempeTy ()))
  unifyMatch _ []                                                              = Right mempty
  unifyMatch um ((ty@(TyBuiltin _ b0), ty'@(TyBuiltin _ b1)):tys) | b0 == b1   = unifyPrep um tys
                                                                  | otherwise  = Left (UnificationFailed () ty ty')
  unifyMatch um ((ty@(TyNamed _ n0), ty'@(TyNamed _ n1)):tys) | n0 == n1       = unifyPrep um tys
                                                              | otherwise      = Left (UnificationFailed () (void ty) (void ty'))
  unifyMatch um ((ty@(TyNamed _ _), TyVar  _ (Name _ (Unique k) _)):tys)       = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch um ((TyVar _ (Name _ (Unique k) _), ty@(TyNamed _ _)):tys)        = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch um ((ty@TyBuiltin{}, TyVar  _ (Name _ (Unique k) _)):tys)         = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch um ((TyVar _ (Name _ (Unique k) _), ty@(TyBuiltin _ _)):tys)      = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch _ ((ty@TyBuiltin{}, ty'@TyNamed{}):_)                             = Left (UnificationFailed () ty ty')
  unifyMatch _ ((ty@TyNamed{}, ty'@TyBuiltin{}):_)                             = Left (UnificationFailed () ty ty')
  unifyMatch _ ((ty@TyBuiltin{}, ty'@TyApp{}):_)                               = Left (UnificationFailed () ty ty')
  unifyMatch _ ((ty@TyNamed{}, ty'@TyApp{}):_)                                 = Left (UnificationFailed () ty ty')
  unifyMatch _ ((ty@TyApp{}, ty'@TyBuiltin{}):_)                               = Left (UnificationFailed () ty ty')
  unifyMatch um ((TyVar _ (Name _ (Unique k) _), ty@TyApp{}):tys)              = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch um ((ty@TyApp{}, TyVar  _ (Name _ (Unique k) _)):tys)             = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  unifyMatch um ((TyApp _ ty ty', TyApp _ ty'' ty'''):tys)                     = unifyPrep um ((ty, ty'') : (ty', ty''') : tys)
  unifyMatch _ ((ty@TyApp{}, ty'@TyNamed{}):_)                                 = Left (UnificationFailed () (void ty) (void ty'))
  unifyMatch um ((TyVar _ n@(Name _ (Unique k) _), ty@(TyVar _ n')):tys)
      | n == n' = unifyPrep um tys -- a type variable is always equal to itself, don't bother inserting this!
      | otherwise = IM.insert k ty <$> unifyPrep (IM.insert k ty um) tys
  
  unify :: [(KempeTy (), KempeTy ())] -> Either (Error ()) (IM.IntMap (KempeTy ()))
  unify = unifyPrep IM.empty
  
  unifyM :: S.Set (KempeTy (), KempeTy ()) -> TypeM () (IM.IntMap (KempeTy ()))
  unifyM s =
      case {-# SCC "unify" #-} unify (S.toList s) of
          Right x  -> pure x
          Left err -> throwError err
  
  -- TODO: take constructor types as an argument?..
  runTypeM :: Int -- ^ For renamer
           -> TypeM a x -> Either (Error a) (x, Int)
  runTypeM maxInt = fmap (second maxU) .
      flip runStateT (TyState maxInt mempty mempty mempty mempty S.empty)
  
  typeOfBuiltin :: BuiltinFn -> TypeM () (StackType ())
  typeOfBuiltin Drop = do
      aN <- dummyName "a"
      pure $ StackType [TyVar () aN] []
  typeOfBuiltin Swap = do
      aN <- dummyName "a"
      bN <- dummyName "b"
      pure $ StackType [TyVar () aN, TyVar () bN] [TyVar () bN, TyVar () aN]
  typeOfBuiltin Dup = do
      aN <- dummyName "a"
      pure $ StackType [TyVar () aN] [TyVar () aN, TyVar () aN]
  typeOfBuiltin IntEq      = pure intRel
  typeOfBuiltin IntLeq     = pure intRel
  typeOfBuiltin IntLt      = pure intRel
  typeOfBuiltin IntMod     = pure intBinOp
  typeOfBuiltin IntDiv     = pure intBinOp
  typeOfBuiltin IntPlus    = pure intBinOp
  typeOfBuiltin IntTimes   = pure intBinOp
  typeOfBuiltin IntMinus   = pure intBinOp
  typeOfBuiltin IntShiftR  = pure intShift
  typeOfBuiltin IntShiftL  = pure intBinOp
  typeOfBuiltin IntXor     = pure intBinOp
  typeOfBuiltin WordXor    = pure wordBinOp
  typeOfBuiltin WordPlus   = pure wordBinOp
  typeOfBuiltin WordTimes  = pure wordBinOp
  typeOfBuiltin WordShiftR = pure wordShift
  typeOfBuiltin WordShiftL = pure wordShift
  typeOfBuiltin IntGeq     = pure intRel
  typeOfBuiltin IntNeq     = pure intRel
  typeOfBuiltin IntGt      = pure intRel
  typeOfBuiltin WordMinus  = pure wordBinOp
  typeOfBuiltin WordDiv    = pure wordBinOp
  typeOfBuiltin WordMod    = pure wordBinOp
  typeOfBuiltin And        = pure boolOp
  typeOfBuiltin Or         = pure boolOp
  typeOfBuiltin Xor        = pure boolOp
  typeOfBuiltin IntNeg     = pure $ StackType [TyBuiltin () TyInt] [TyBuiltin () TyInt]
  typeOfBuiltin Popcount   = pure $ StackType [TyBuiltin () TyWord] [TyBuiltin () TyInt]
  
  boolOp :: StackType ()
  boolOp = StackType [TyBuiltin () TyBool, TyBuiltin () TyBool] [TyBuiltin () TyBool]
  
  intRel :: StackType ()
  intRel = StackType [TyBuiltin () TyInt, TyBuiltin () TyInt] [TyBuiltin () TyBool]
  
  intBinOp :: StackType ()
  intBinOp = StackType [TyBuiltin () TyInt, TyBuiltin () TyInt] [TyBuiltin () TyInt]
  
  intShift :: StackType ()
  intShift = StackType [TyBuiltin () TyInt, TyBuiltin () TyInt] [TyBuiltin () TyInt]
  
  wordBinOp :: StackType ()
  wordBinOp = StackType [TyBuiltin () TyWord, TyBuiltin () TyWord] [TyBuiltin () TyWord]
  
  wordShift :: StackType ()
  wordShift = StackType [TyBuiltin () TyWord, TyBuiltin () TyWord] [TyBuiltin () TyWord]
  
  tyLookup :: Name a -> TypeM a (StackType a)
  tyLookup n@(Name _ (Unique i) l) = do
      st <- gets tyEnv
      case IM.lookup i st of
          Just ty -> pure ty
          Nothing -> throwError $ PoorScope l n
  
  consLookup :: TyName a -> TypeM a (StackType a)
  consLookup tn@(Name _ (Unique i) l) = do
      st <- gets constructorTypes
      case IM.lookup i st of
          Just ty -> pure ty
          Nothing -> throwError $ PoorScope l tn
  
  -- expandType 1
  dipify :: StackType () -> TypeM () (StackType ())
  dipify (StackType is os) = do
      n <- dummyName "a"
      pure $ StackType (is ++ [TyVar () n]) (os ++ [TyVar () n])
  
  tyLeaf :: (Pattern b a, [Atom b a]) -> TypeM () (StackType ())
  tyLeaf (p, as) = do
      tyP <- tyPattern p
      tyA <- tyAtoms as
      catTypes tyP tyA
  
  assignCase :: (Pattern b a, [Atom b a]) -> TypeM () (StackType (), Pattern (StackType ()) (StackType ()), [Atom (StackType ()) (StackType ())])
  assignCase (p, as) = do
      (tyP, p') <- assignPattern p
      (as', tyA) <- assignAtoms as
      (,,) <$> catTypes tyP tyA <*> pure p' <*> pure as'
  
  tyAtom :: Atom b a -> TypeM () (StackType ())
  tyAtom (AtBuiltin _ b) = typeOfBuiltin b
  tyAtom BoolLit{}       = pure $ StackType [] [TyBuiltin () TyBool]
  tyAtom IntLit{}        = pure $ StackType [] [TyBuiltin () TyInt]
  tyAtom Int8Lit{}       = pure $ StackType [] [TyBuiltin () TyInt8 ]
  tyAtom WordLit{}       = pure $ StackType [] [TyBuiltin () TyWord]
  tyAtom (AtName _ n)    = renameStack =<< tyLookup (void n)
  tyAtom (Dip _ as)      = dipify =<< tyAtoms as
  tyAtom (AtCons _ tn)   = renameStack =<< consLookup (void tn)
  tyAtom (If _ as as')   = do
      tys <- tyAtoms as
      tys' <- tyAtoms as'
      (StackType ins out) <- mergeStackTypes tys tys'
      pure $ StackType (ins ++ [TyBuiltin () TyBool]) out
  tyAtom (Case _ ls) = do
      tyLs <- traverse tyLeaf ls
      -- TODO: one-pass fold?
      mergeMany tyLs
  
  assignAtom :: Atom b a -> TypeM () (StackType (), Atom (StackType ()) (StackType ()))
  assignAtom (AtBuiltin _ b) = do { ty <- typeOfBuiltin b ; pure (ty, AtBuiltin ty b) }
  assignAtom (BoolLit _ b)   =
      let sTy = StackType [] [TyBuiltin () TyBool]
          in pure (sTy, BoolLit sTy b)
  assignAtom (IntLit _ i)    =
      let sTy = StackType [] [TyBuiltin () TyInt]
          in pure (sTy, IntLit sTy i)
  assignAtom (Int8Lit _ i)    =
      let sTy = StackType [] [TyBuiltin () TyInt8]
          in pure (sTy, Int8Lit sTy i)
  assignAtom (WordLit _ u)    =
      let sTy = StackType [] [TyBuiltin () TyWord]
          in pure (sTy, WordLit sTy u)
  assignAtom (AtName _ n) = do
      sTy <- renameStack =<< tyLookup (void n)
      pure (sTy, AtName sTy (n $> sTy))
  assignAtom (AtCons _ tn) = do
      sTy <- renameStack =<< consLookup (void tn)
      pure (sTy, AtCons sTy (tn $> sTy))
  assignAtom (Dip _ as) = do { (as', ty) <- assignAtoms as ; tyDipped <- dipify ty ; pure (tyDipped, Dip tyDipped as') }
  assignAtom (If _ as0 as1) = do
      (as0', tys) <- assignAtoms as0
      (as1', tys') <- assignAtoms as1
      (StackType ins out) <- mergeStackTypes tys tys'
      let resType = StackType (ins ++ [TyBuiltin () TyBool]) out
      pure (resType, If resType as0' as1')
  assignAtom (Case _ ls) = do
      lRes <- traverse assignCase ls
      resType <- mergeMany (fst3 <$> lRes)
      let newLeaves = fmap dropFst lRes
      pure (resType, Case resType newLeaves)
      where dropFst (_, y, z) = (y, z)
  
  assignAtoms :: [Atom b a] -> TypeM () ([Atom (StackType ()) (StackType ())], StackType ())
  assignAtoms []  = pure ([], emptyStackType)
  assignAtoms [a] = do
      (ty, a') <- assignAtom a
      pure ([a'], ty)
  assignAtoms (a:as) = do
      (ty, a') <- assignAtom a
      (as', ty') <- assignAtoms as
      (a':as' ,) <$> catTypes ty ty'
  
  tyAtoms :: [Atom b a] -> TypeM () (StackType ())
  tyAtoms [] = pure emptyStackType
  tyAtoms [a] = tyAtom a
  tyAtoms (a:as) = do
      ty <- tyAtom a
      tys <- tyAtoms as
      catTypes ty tys
  
  -- from size,
  mkHKT :: Int -> Kind
  mkHKT 0 = Star
  mkHKT i = TyCons (mkHKT $ i - 1) Star
  
  tyInsertLeaf :: Name b -- ^ type being declared
               -> S.Set (Name b) -> (TyName a, [KempeTy b]) -> TypeM () ()
  tyInsertLeaf n@(Name _ (Unique k) _) vars (Name _ (Unique i) _, ins) | S.null vars =
      modifying constructorTypesLens (IM.insert i (voidStackType $ StackType ins [TyNamed undefined n])) *>
      modifying kindEnvLens (IM.insert k Star)
                                                 | otherwise =
      let ty = voidStackType $ StackType ins [app (TyNamed undefined n) (S.toList vars)] in
      modifying constructorTypesLens (IM.insert i ty) *>
      modifying kindEnvLens (IM.insert k (mkHKT $ S.size vars))
  
  assignTyLeaf :: Name b
               -> S.Set (Name b)
               -> (TyName a, [KempeTy b])
               -> TypeM () (TyName (StackType ()), [KempeTy ()])
  assignTyLeaf n@(Name _ (Unique k) _) vars (tn@(Name _ (Unique i) _), ins) | S.null vars =
      let ty = voidStackType $ StackType ins [TyNamed undefined n] in
      modifying constructorTypesLens (IM.insert i ty) *>
      modifying kindEnvLens (IM.insert k Star) $>
      (tn $> ty, fmap void ins)
                                                 | otherwise =
      let ty = voidStackType $ StackType ins [app (TyNamed undefined n) (S.toList vars)] in
      modifying constructorTypesLens (IM.insert i ty) *>
      modifying kindEnvLens (IM.insert k (mkHKT $ S.size vars)) $>
      (tn $> ty, fmap void ins)
  
  app :: KempeTy a -> [Name a] -> KempeTy a
  app = foldr (\n ty -> TyApp undefined ty (TyVar undefined n))
  
  kindLookup :: TyName a -> TypeM a Kind
  kindLookup n@(Name _ (Unique i) l) = do
      st <- gets kindEnv
      case IM.lookup i st of
          Just k  -> pure k
          Nothing -> throwError $ PoorScope l n
  
  kindOf :: KempeTy a -> TypeM a Kind
  kindOf TyBuiltin{}        = pure Star
  kindOf (TyNamed _ tn)     = kindLookup tn
  kindOf TyVar{}            = pure Star
  kindOf tyErr@(TyApp l ty ty') = do
      k <- kindOf ty
      k' <- kindOf ty'
      case k of
          TyCons k'' k''' -> unless (k' == k''') (throwError (IllKinded l tyErr)) $> k''
          _               -> throwError (IllKinded l tyErr)
  
  assignDecl :: KempeDecl a c b -> TypeM () (KempeDecl () (StackType ()) (StackType ()))
  assignDecl (TyDecl _ tn ns ls) = TyDecl () (void tn) (void <$> ns) <$> traverse (assignTyLeaf tn (S.fromList ns)) ls
  assignDecl (FunDecl _ n ins os a) = do
      traverse_ kindOf (void <$> ins ++ os)
      sig <- renameStack $ voidStackType $ StackType ins os
      (as, inferred) <- assignAtoms a
      reconcile <- mergeStackTypes sig inferred
      when (inferred `lessGeneral` sig) $
          throwError $ LessGeneral () sig inferred
      pure $ FunDecl reconcile (n $> reconcile) (void <$> ins) (void <$> os) as
  assignDecl (ExtFnDecl _ n ins os cn) = do
      traverse_ kindOf (void <$> ins ++ os)
      unless (length os <= 1) $
          throwError $ InvalidCImport () (void n)
      let sig = voidStackType $ StackType ins os
      pure $ ExtFnDecl sig (n $> sig) (void <$> ins) (void <$> os) cn
  assignDecl (Export _ abi n) = do
      ty@(StackType _ os) <- tyLookup (void n)
      unless (abi == Kabi || length os <= 1) $
          throwError $ InvalidCExport () (void n)
      Export ty abi <$> assignName n
  
  -- don't need to rename cuz it's only for exports (in theory)
  assignName :: Name a -> TypeM () (Name (StackType ()))
  assignName n = do { ty <- tyLookup (void n) ; pure (n $> ty) }
  
  tyHeader :: KempeDecl a c b -> TypeM () ()
  tyHeader Export{} = pure ()
  tyHeader (FunDecl _ (Name _ (Unique i) _) ins out _) = do
      let sig = voidStackType $ StackType ins out
      modifying tyEnvLens (IM.insert i sig)
  tyHeader (ExtFnDecl _ n@(Name _ (Unique i) _) ins os _) = do
      unless (length os <= 1) $
          throwError $ InvalidCImport () (void n)
      unless (null $ freeVars (ins ++ os)) $
          throwError $ TyVarExt () (void n)
      let sig = voidStackType $ StackType ins os -- no free variables allowed in c functions
      modifying tyEnvLens (IM.insert i sig)
  tyHeader TyDecl{} = pure ()
  
  type Vars a = IM.IntMap (Name a)
  
  -- TODO: do we want strict or lazy?
  type EqState a = State (Vars a)
  
  -- need to check stack types are less general up to "alpha-equivalence"
  -- (implicit forall with every new var! in a stack type)
  --
  -- Basically the inferred type has to check against the type in the signature.
  -- Which I'm not sure this does; it should be better-founded with an
  -- explanation of why it works (I'm not sure it works)
  lessGeneral :: StackType a -- ^ Inferred type
              -> StackType a -- ^ Type from signature
              -> Bool
  lessGeneral (StackType is os) (StackType is' os') =
      flip evalState mempty $
          if il > il' || ol > ol'
              then (||) <$> lessGenerals trimIs is' <*> lessGenerals trimOs os'
              else (||) <$> lessGenerals is trimIs' <*> lessGenerals os trimOs'
      where il = length is
            il' = length is'
            ol = length os
            ol' = length os'
            trimIs = drop (il-il') is
            trimIs' = drop (il'-il) is'
            trimOs = drop (ol-ol') os
            trimOs' = drop (ol'-ol) os'
  
            lessGeneralAtom :: KempeTy a -> KempeTy a -> EqState a Bool
            lessGeneralAtom TyBuiltin{} TyVar{}                   = pure True
            lessGeneralAtom TyApp{} TyVar{}                       = pure True
            lessGeneralAtom (TyApp _ ty ty') (TyApp _ ty'' ty''') = (||) <$> lessGeneralAtom ty ty'' <*> lessGeneralAtom ty' ty''' -- lazy pattern match?
            lessGeneralAtom _ _                                   = pure False
            lessGenerals :: [KempeTy a] -> [KempeTy a] -> EqState a Bool
            lessGenerals [] []                                 = pure False
            lessGenerals ((TyVar _ n):tys) ((TyVar _ n'):tys') = do
                  st <- get
                  let i = unUnique $ unique n
                  case IM.lookup i st of
                      Nothing ->
                          modify (IM.insert i n') *>
                          lessGenerals tys tys' -- we can skip checking ty `lessGeneral` ty' at the first site
                      Just n'' ->
                          (n'' /= n' ||) <$> lessGenerals tys tys'
            lessGenerals (ty:tys) (ty':tys')                  = (||) <$> lessGeneralAtom ty ty' <*> lessGenerals tys tys'
  
  tyInsert :: KempeDecl a c b -> TypeM () ()
  tyInsert (TyDecl _ tn ns ls) = traverse_ (tyInsertLeaf tn (S.fromList ns)) ls
  tyInsert (FunDecl _ _ ins out as) = do
      traverse_ kindOf (void <$> ins ++ out) -- FIXME: this gives sketchy results?
      sig <- renameStack $ voidStackType $ StackType ins out
      inferred <- tyAtoms as
      _ <- mergeStackTypes sig inferred
      when (inferred `lessGeneral` sig) $
          throwError $ LessGeneral () sig inferred
  tyInsert (ExtFnDecl _ _ ins outs _) = traverse_ kindOf (void <$> ins ++ outs)
  tyInsert Export{} = pure ()
  
  tyModule :: Declarations a c b -> TypeM () ()
  tyModule m = traverse_ tyHeader m *> traverse_ tyInsert m
  
  checkModule :: Declarations a c b -> TypeM () ()
  checkModule m = tyModule m <* (unifyM =<< gets constraints)
  
  assignModule :: Declarations a c b -> TypeM () (Declarations () (StackType ()) (StackType ()))
  assignModule m = {-# SCC "assignModule" #-} do
      {-# SCC "tyHeader" #-} traverse_ tyHeader m
      m' <- traverse assignDecl m
      backNames <- unifyM =<< gets constraints
      pure (fmap (bimap .$ substConstraintsStack backNames) m')
  
  -- Make sure you don't have cycles in the renames map!
  replaceUnique :: Unique -> TypeM a Unique
  replaceUnique u@(Unique i) = do
      rSt <- gets renames
      case IM.lookup i rSt of
          Nothing -> pure u
          Just j  -> replaceUnique (Unique j)
  
  renameIn :: KempeTy a -> TypeM a (KempeTy a)
  renameIn b@TyBuiltin{}    = pure b
  renameIn n@TyNamed{}      = pure n
  renameIn (TyApp l ty ty') = TyApp l <$> renameIn ty <*> renameIn ty'
  renameIn (TyVar l (Name t u l')) = do
      u' <- replaceUnique u
      pure $ TyVar l (Name t u' l')
  
  -- has to use the max-iest maximum so we can't use withState
  withTyState :: (TyState a -> TyState a) -> TypeM a x -> TypeM a x
  withTyState modSt act = do
      preSt <- get
      modify modSt
      res <- act
      postMax <- gets maxU
      put preSt
      maxULens .= postMax
      pure res
  
  withName :: Name a -> TypeM a (Name a, TyState a -> TyState a)
  withName (Name t (Unique i) l) = do
      m <- gets maxU
      let newUniq = m+1
      maxULens .= newUniq
      pure (Name t (Unique newUniq) l, over renamesLens (IM.insert i (m+1)))
  
  -- freshen the names in a stack so there aren't overlaps in quanitified variables
  renameStack :: StackType a -> TypeM a (StackType a)
  renameStack (StackType ins outs) = do
      newQs <- traverse withName (S.toList $ freeVars ins <> freeVars outs)
      let (_, localRenames) = unzip newQs
          newBinds = thread localRenames
      withTyState newBinds $
          StackType <$> traverse renameIn ins <*> traverse renameIn outs
  
  mergeStackTypes :: StackType () -> StackType () -> TypeM () (StackType ())
  mergeStackTypes st0@(StackType i0 o0) st1@(StackType i1 o1) = do
      let li0 = length i0
          li1 = length i1
          toExpand = max (abs (li0 - li1)) (abs (length o0 - length o1))
  
      (StackType ins os) <- (if li0 < li1 then expandType toExpand else pure) st0
      (StackType ins' os') <- (if li1 < li0 then expandType toExpand else pure) st1
  
      when ((length ins /= length ins') || (length os /= length os')) $
          throwError $ MismatchedLengths () st0 st1
  
      zipWithM_ pushConstraint ins ins'
      zipWithM_ pushConstraint os os'
  
      pure $ StackType ins os
  
  tyPattern :: Pattern b a -> TypeM () (StackType ())
  tyPattern PatternWildcard{} = do
      aN <- dummyName "a"
      pure $ StackType [TyVar () aN] []
  tyPattern PatternInt{} = pure $ StackType [TyBuiltin () TyInt] []
  tyPattern PatternBool{} = pure $ StackType [TyBuiltin () TyBool] []
  tyPattern (PatternCons _ tn) = renameStack . flipStackType =<< consLookup (void tn)
  
  assignPattern :: Pattern b a -> TypeM () (StackType (), Pattern (StackType ()) (StackType ()))
  assignPattern (PatternInt _ i) =
      let sTy = StackType [TyBuiltin () TyInt] []
          in pure (sTy, PatternInt sTy i)
  assignPattern (PatternBool _ i) =
      let sTy = StackType [TyBuiltin () TyBool] []
          in pure (sTy, PatternBool sTy i)
  assignPattern (PatternCons _ tn) = do { ty <- renameStack . flipStackType =<< consLookup (void tn) ; pure (ty, PatternCons ty (tn $> ty)) }
  assignPattern PatternWildcard{} = do
      aN <- dummyName "a"
      let resType = StackType [TyVar () aN] []
      pure (resType, PatternWildcard resType)
  
  mergeMany :: NonEmpty (StackType ()) -> TypeM () (StackType ())
  mergeMany (t :| ts) = foldM mergeStackTypes t ts
  
  -- assumes they have been renamed...
  pushConstraint :: Ord a => KempeTy a -> KempeTy a -> TypeM a ()
  pushConstraint ty ty' =
      modifying constraintsLens (S.insert (ty, ty'))
  
  expandType :: Int -> StackType () -> TypeM () (StackType ())
  expandType n (StackType i o) = do
      newVars <- replicateM n (dummyName "a")
      let newTy = TyVar () <$> newVars
      pure $ StackType (newTy ++ i) (newTy ++ o)
  
  substConstraints :: IM.IntMap (KempeTy a) -> KempeTy a -> KempeTy a
  substConstraints _ ty@TyNamed{}                         = ty
  substConstraints _ ty@TyBuiltin{}                       = ty
  substConstraints tys ty@(TyVar _ (Name _ (Unique k) _)) =
      case IM.lookup k tys of
          Just ty'@TyVar{}       -> substConstraints (IM.delete k tys) ty' -- TODO: this is to prevent cyclic lookups: is it right?
          Just (TyApp l ty0 ty1) -> let tys' = IM.delete k tys in TyApp l (substConstraints tys' ty0) (substConstraints tys' ty1)
          Just ty'               -> ty'
          Nothing                -> ty
  substConstraints tys (TyApp l ty ty')                   =
      TyApp l (substConstraints tys ty) (substConstraints tys ty')
  
  substConstraintsStack :: IM.IntMap (KempeTy a) -> StackType a -> StackType a
  substConstraintsStack tys (StackType is os) = {-# SCC "substConstraintsStack" #-}
      let is' = substConstraints tys <$> is
          os' = substConstraints tys <$> os
          in StackType is' os'
  
  -- do renaming before this
  -- | Given @x@ and @y@, return the 'StackType' of @x y@
  catTypes :: StackType () -- ^ @x@
           -> StackType () -- ^ @y@
           -> TypeM () (StackType ())
  catTypes st0@(StackType _ osX) (StackType insY osY) = do
      let lY = length insY
          lDiff = lY - length osX
  
      -- all of the "ins" of y have to come from x, so we expand x as needed
      (StackType insX osX') <- if lDiff > 0
          then expandType lDiff st0
          else pure st0
  
      -- zip the last (length insY) of osX' with insY
      zipWithM_ pushConstraint (drop (length osX' - lY) osX') insY -- TODO splitAt
  
      pure $ StackType insX (take (length osX' - lY) osX' ++ osY)
  