content: {-# LANGUAGE FlexibleContexts  #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TupleSections     #-}
  
  -- | This module is kind of half-assed. I don't have any references and it
  -- depends on the inliner.
  module Kempe.Monomorphize ( closedModule
                            , MonoM
                            , runMonoM
                            , flattenModule
                            , tryMono
                            , ConsAnn (..)
                            -- * Benchmark
                            , closure
                            , mkModuleMap
                            ) where
  
  import           Control.Arrow              ((&&&))
  import           Control.Monad              ((<=<))
  import           Control.Monad.Except       (MonadError, throwError)
  import           Control.Monad.State.Strict (StateT, gets, runStateT)
  import           Data.Bifunctor             (second)
  import           Data.Containers.ListUtils  (nubOrd)
  import           Data.Foldable              (traverse_)
  import           Data.Function              (on)
  import           Data.Functor               (($>))
  import qualified Data.IntMap                as IM
  import           Data.List                  (elemIndex, find, groupBy, partition)
  import qualified Data.Map                   as M
  import           Data.Maybe                 (fromMaybe, mapMaybe)
  import qualified Data.Set                   as S
  import qualified Data.Text                  as T
  import           Data.Tuple                 (swap)
  import           Data.Tuple.Ext             (fst3, snd3, thd3)
  import           Kempe.AST
  import           Kempe.AST.Size
  import           Kempe.Error
  import           Kempe.Name
  import           Kempe.Unique
  import           Lens.Micro                 (Lens')
  import           Lens.Micro.Mtl             (modifying)
  import           Prettyprinter              (Doc, Pretty, vsep)
  import           Prettyprinter.Debug
  
  -- | New function names, keyed by name + specialized type
  --
  -- also max state threaded through.
  data RenameEnv = RenameEnv { maxState :: Int
                             , fnEnv    :: M.Map (Unique, StackType ()) Unique
                             , consEnv  :: M.Map (Unique, StackType ()) (Unique, ConsAnn MonoStackType)
                             , szEnv    :: SizeEnv
                             }
  
  type MonoM = StateT RenameEnv (Either (Error ()))
  
  prettyDumpBinds :: (Pretty b, Pretty k) => M.Map k b -> Doc a
  prettyDumpBinds b = vsep (prettyBind <$> M.toList b)
  
  maxStateLens :: Lens' RenameEnv Int
  maxStateLens f s = fmap (\x -> s { maxState = x }) (f (maxState s))
  
  consEnvLens :: Lens' RenameEnv (M.Map (Unique, StackType ()) (Unique, ConsAnn MonoStackType))
  consEnvLens f s = fmap (\x -> s { consEnv = x }) (f (consEnv s))
  
  fnEnvLens :: Lens' RenameEnv (M.Map (Unique, StackType ()) Unique)
  fnEnvLens f s = fmap (\x -> s { fnEnv = x }) (f (fnEnv s))
  
  szEnvLens :: Lens' RenameEnv SizeEnv
  szEnvLens f s = fmap (\x -> s { szEnv = x }) (f (szEnv s))
  
  runMonoM :: Int -> MonoM a -> Either (Error ()) (a, (Int, SizeEnv))
  runMonoM maxI = fmap (second (maxState &&& szEnv)) . flip runStateT (RenameEnv maxI mempty mempty mempty)
  
  freshName :: T.Text -> a -> MonoM (Name a)
  freshName n ty = do
      pSt <- gets maxState
      Name n (Unique $ pSt + 1) ty
          <$ modifying maxStateLens (+1)
  
  tryMono :: MonadError (Error ()) m => StackType () -> m MonoStackType
  tryMono (StackType is os) | S.null (freeVars (is ++ os)) = pure (is, os)
                            | otherwise = throwError $ MonoFailed ()
  
  -- | A 'ModuleMap' is a map which retrives the 'KempeDecl' associated with
  -- a given 'Name'
  type ModuleMap a c b = IM.IntMap (KempeDecl a c b)
  
  mkModuleMap :: Declarations a c b -> ModuleMap a c b
  mkModuleMap = IM.fromList . concatMap toInt where
      toInt d@(FunDecl _ (Name _ (Unique i) _) _ _ _)   = [(i, d)]
      toInt d@(ExtFnDecl _ (Name _ (Unique i) _) _ _ _) = [(i, d)]
      toInt d@(TyDecl _ _ _ ds)                         =
          let us = unUnique . unique . fst <$> ds
              in (, d) <$> us
      toInt _                                           = []
  
  squishTypeName :: BuiltinTy -> T.Text
  squishTypeName TyInt  = "int"
  squishTypeName TyBool = "bool"
  squishTypeName TyWord = "word"
  squishTypeName TyInt8 = "int8"
  
  squishType :: KempeTy a -> T.Text
  squishType (TyBuiltin _ b)          = squishTypeName b
  squishType (TyNamed _ (Name t _ _)) = T.toLower t
  squishType TyVar{}                  = error "not meant to be monomorphized!"
  squishType (TyApp _ ty ty')         = squishType ty <> squishType ty'
  
  squishMonoStackType :: MonoStackType -> T.Text
  squishMonoStackType (is, os) = foldMap squishType is <> "TT" <> foldMap squishType os
  
  renamePattern :: Pattern (StackType ()) (StackType ()) -> MonoM (Pattern (ConsAnn MonoStackType) (StackType ()))
  renamePattern (PatternInt ty i)    = pure $ PatternInt ty i
  renamePattern (PatternWildcard ty) = pure $ PatternWildcard ty
  renamePattern (PatternBool ty b)   = pure $ PatternBool ty b
  renamePattern (PatternCons ty (Name t u _)) = do
      cSt <- gets consEnv
      let (u', ann) = M.findWithDefault (error "Internal error? unfound constructor") (u, flipStackType ty) cSt
          ann' = swap <$> ann
      pure $ PatternCons ann' (Name t u' ann')
  
  renameCase :: (Pattern (StackType ()) (StackType ()), [Atom (StackType ()) (StackType ())]) -> MonoM (Pattern (ConsAnn MonoStackType) (StackType ()), [Atom (ConsAnn MonoStackType) (StackType ())])
  renameCase (p, as) = (,) <$> renamePattern p <*> traverse renameAtom as
  
  renameAtom :: Atom (StackType ()) (StackType ()) -> MonoM (Atom (ConsAnn MonoStackType) (StackType ()))
  renameAtom (AtBuiltin ty b)         = pure $ AtBuiltin ty b
  renameAtom (If ty as as')           = If ty <$> traverse renameAtom as <*> traverse renameAtom as'
  renameAtom (IntLit ty i)            = pure $ IntLit ty i
  renameAtom (Int8Lit ty i)           = pure $ Int8Lit ty i
  renameAtom (WordLit ty w)           = pure $ WordLit ty w
  renameAtom (BoolLit ty b)           = pure $ BoolLit ty b
  renameAtom (Dip ty as)              = Dip ty <$> traverse renameAtom as
  renameAtom (AtName ty (Name t u l)) = do
      mSt <- gets fnEnv
      let u' = M.findWithDefault u (u, ty) mSt
      pure $ AtName ty (Name t u' l)
  renameAtom (Case ty ls)             = Case ty <$> traverse renameCase ls
  renameAtom (AtCons ty (Name t u _)) = do
      cSt <- gets consEnv
      let (u', ann) = M.findWithDefault (error "Internal error? unfound constructor") (u, ty) cSt
      pure $ AtCons ann (Name t u' ann)
  
  renameDecl :: KempeDecl () (StackType ()) (StackType ()) -> MonoM (KempeDecl () (ConsAnn MonoStackType) (StackType ()))
  renameDecl (FunDecl l n is os as) = FunDecl l n is os <$> traverse renameAtom as
  renameDecl (Export ty abi (Name t u l)) = do
      mSt <- gets fnEnv
      let u' = M.findWithDefault (error "Shouldn't happen; might be user error or internal error") (u, ty) mSt
      pure $ Export ty abi (Name t u' l)
  renameDecl (ExtFnDecl l n tys tys' b) = pure $ ExtFnDecl l n tys tys' b
  renameDecl (TyDecl l n vars ls)       = pure $ TyDecl l n vars ls
  
  -- | Call 'closedModule' and perform any necessary renamings
  flattenModule :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (ConsAnn MonoStackType) (StackType ()))
  flattenModule = renameMonoM <=< closedModule
  
  -- | To be called after 'closedModule'
  renameMonoM :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (ConsAnn MonoStackType) (StackType ()))
  renameMonoM = traverse renameDecl
  
  -- | Filter so that only the 'KempeDecl's necessary for exports are there, and
  -- fan out top-level functions into all necessary specializations.
  --
  -- This will throw an exception on ill-typed programs.
  --
  -- The 'Module' returned will have to be renamed.
  closedModule :: Declarations () (StackType ()) (StackType ()) -> MonoM (Declarations () (StackType ()) (StackType ()))
  closedModule m = addExports <$> do
      { fn' <- traverse (uncurry specializeDecl . drop1) fnDecls
      ; traverse_ insTyDecl $ nubOrd (snd3 <$> tyDecls)
      ; ty' <- specializeTyDecls tyDecls
      ; pure (ty' ++ fn')
      }
      where addExports = (++ exportsOnly m)
            key = mkModuleMap m
            roots = S.toList $ closure (m, key)
            gatherDecl (n@(Name _ (Unique i) _), ty) = -- TODO: findWithDefault?
              case IM.lookup i key of
                  Just decl -> (n, decl, ty)
                  Nothing   -> error "Internal error! module map should contain all names."
            rootDecl = gatherDecl <$> roots -- FIXME: two-steps away, the roots are not monomorphized! So it tries to create specialized declarations of type a b -- a b a &c.
            drop1 ~(_, y, z) = (y, z)
            (tyDecls, fnDecls) = partition (isTyDecl . snd3) rootDecl
            isTyDecl TyDecl{} = True
            isTyDecl _        = False
  
  -- group specializations by type name?
  specializeTyDecls :: [(TyName (StackType ()), KempeDecl () (StackType ()) (StackType ()), StackType ())] -> MonoM [KempeDecl () (StackType ()) (StackType ())]
  specializeTyDecls ds = traverse (uncurry mkTyDecl) processed
      where toMerge = groupBy ((==) `on` snd3) ds
            processed = fmap process toMerge
            process tyDs@((_, x, _):_) = (x, zip (fst3 <$> tyDs) (thd3 <$> tyDs))
            process []                 = error "Empty group!"
  
  isTyVar :: KempeTy a -> Bool
  isTyVar TyVar{} = True
  isTyVar _       = False
  
  extrNames :: KempeTy a -> Name a
  extrNames (TyVar _ n) = n
  extrNames _           = error "Internal error!"
  
  sizeLeaf :: [Name a] -- ^ Type variables as declared
           -> [KempeTy a]
           -> MonoM Size
  sizeLeaf fv tys = do
      { let (tvs, conc) = partition isTyVar tys
      ; pad <- sizeStack <$> gets szEnv <*> pure conc
      ; let tvPrecompose = fmap (forVar . extrNames) tvs
      ; let tvComposed = foldr compose (const pad) tvPrecompose
      ; pure tvComposed
      }
    where
      findIx x = fromMaybe (error "Internal error: can't find index of type variable.") $ elemIndex x fv
      forVar n =
          let i = findIx n
              in (!! i)
      compose sz sz' = \tys' -> sz tys' + sz' tys'
  
  insTyDecl :: KempeDecl a c b -> MonoM ()
  insTyDecl (TyDecl _ (Name _ (Unique k) _) fv leaves) = do
      leafSizes <- traverse (sizeLeaf fv.snd) leaves
      let consSz = \tys -> 1 + maximum (($tys) <$> leafSizes) -- for the tag
      modifying szEnvLens (IM.insert k consSz)
  insTyDecl _ = error "Shouldn't happen."
  
  mkTyDecl :: KempeDecl () (StackType ()) (StackType ()) -> [(TyName (StackType ()), StackType ())] -> MonoM (KempeDecl () (StackType ()) (StackType ()))
  mkTyDecl (TyDecl _ tn ns preConstrs) constrs = do
      env <- gets szEnv
      renCons <- traverse (\(tn', ty) -> do { ty'@(is, _) <- tryMono ty ; (, is) <$> renamedCons (tn' $> ty') ty' (ConsAnn (szType env ty') (getTag tn')) }) constrs
      pure $ TyDecl () tn ns renCons
      where indexAt p xs = fst $ fromMaybe (error "Internal error.") $ find (\(_, x) -> p x) (zip [0..] xs)
            getTag (Name _ u _) = indexAt (== u) preIxes
            preIxes = fmap (unique . fst) preConstrs
            szType env (_, [o]) = size' env o
            szType _ _          = error "Internal error: ill-typed constructor."
  mkTyDecl _ _ = error "Shouldn't happen."
  
  specializeDecl :: KempeDecl () (StackType ()) (StackType ()) -> StackType () -> MonoM (KempeDecl () (StackType ()) (StackType ()))
  specializeDecl (FunDecl _ n _ _ as) sty = do
      (Name t u newStackType@(StackType is os)) <- renamed n =<< tryMono sty
      pure $ FunDecl newStackType (Name t u newStackType) is os as
  specializeDecl (ExtFnDecl l n tys tys' b) _ = pure $ ExtFnDecl l n tys tys' b
  specializeDecl (Export l abi n) _           = pure $ Export l abi n
  specializeDecl TyDecl{} _                   = error "Shouldn't happen."
  -- leave exports and foreign imports alone (have to be monomorphic)
  
  renamedCons :: TyName a -> MonoStackType -> (MonoStackType -> ConsAnn MonoStackType) -> MonoM (TyName (StackType ()))
  renamedCons (Name t i _) sty@(is, os) fAnn = do
      let t' = t <> squishMonoStackType sty
      (Name _ j _) <- freshName t' sty
      let newStackType = StackType is os
          ann = fAnn sty
      modifying consEnvLens (M.insert (i, newStackType) (j, ann))
      pure (Name t' j newStackType)
  
  -- | Insert a specialized rename.
  renamed :: Name a -> MonoStackType -> MonoM (Name (StackType ()))
  renamed (Name t i _) sty@(is, os) = do
      let t' = t <> squishMonoStackType sty
      (Name _ j _) <- freshName t' sty
      let newStackType = StackType is os
      modifying fnEnvLens (M.insert (i, newStackType) j)
      pure (Name t' j newStackType)
  
  closure :: (Declarations a (StackType ()) (StackType ()), ModuleMap a (StackType ()) (StackType ())) -> S.Set (Name (StackType ()), StackType ())
  closure (m, key) = loop roots S.empty
      where roots = S.fromList (exports m)
            loop ns avoid =
              let res = foldMap (step . fst) (ns S.\\ avoid)
                  in if res == ns
                      then res
                      else ns <> loop res (ns <> avoid)
            step (Name _ (Unique i) _) =
              case IM.lookup i key of
                  Just decl -> namesInDecl decl
                  Nothing   -> error "Internal error! module map should contain all names."
  
  namesInDecl :: KempeDecl a (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
  namesInDecl TyDecl{}             = S.empty
  namesInDecl ExtFnDecl{}          = S.empty
  namesInDecl Export{}             = S.empty
  namesInDecl (FunDecl _ _ _ _ as) = foldMap namesInAtom as
  
  namesInAtom :: Atom (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
  namesInAtom AtBuiltin{}                = S.empty
  namesInAtom (If _ as as')              = foldMap namesInAtom as <> foldMap namesInAtom as'
  namesInAtom (Dip _ as)                 = foldMap namesInAtom as
  namesInAtom (AtName _ n@(Name _ _ l))  = S.singleton (n, l)
  namesInAtom (AtCons _ tn@(Name _ _ l)) = S.singleton (tn, l)
  namesInAtom IntLit{}                   = S.empty
  namesInAtom BoolLit{}                  = S.empty
  namesInAtom Int8Lit{}                  = S.empty
  namesInAtom WordLit{}                  = S.empty
  namesInAtom (Case _ as)                = foldMap namesInAtom (foldMap snd as) <> foldMap (namesInPattern . fst) as
  
  namesInPattern :: Pattern (StackType ()) (StackType ()) -> S.Set (Name (StackType ()), StackType ())
  namesInPattern (PatternCons _ tn@(Name _ _ l)) = S.singleton (tn, flipStackType l) -- we have to flipStackType here because the type is the reverse of the constructor that we want to look up
  namesInPattern _                               = S.empty
  
  exports :: Declarations a c b -> [(Name b, b)]
  exports = mapMaybe exportsDecl
  
  exportsOnly :: Declarations a c b -> Declarations a c b
  exportsOnly = mapMaybe getExport where
      getExport d@Export{} = Just d
      getExport _          = Nothing
  
  exportsDecl :: KempeDecl a c b -> Maybe (Name b, b)
  exportsDecl (Export _ _ n@(Name _ _ l)) = Just (n, l)
  exportsDecl _                           = Nothing
  