<commit_before>module Malgo.TypeCheck.Pass where

import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import Data.List (find, head, last, splitAt)
import Data.List.Extra (anySame)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromJust)
import Koriel.Id
import Koriel.MonadUniq
import Koriel.Pretty
import Malgo.Interface (loadInterface, signatureMap, typeDefMap)
import Malgo.Prelude
import Malgo.Rename.RnEnv (RnEnv)
import qualified Malgo.Rename.RnEnv as R
import Malgo.Syntax hiding (Type (..), freevars)
import qualified Malgo.Syntax as S
import Malgo.Syntax.Extension
import Malgo.TypeCheck.TcEnv
import Malgo.TypeRep.Static (Rep (..), Scheme (Forall), TypeDef (..), TypeF, typeConstructor, typeParameters, valueConstructors)
import qualified Malgo.TypeRep.Static as Static
import Malgo.TypeRep.UTerm
import Malgo.UTerm
import Malgo.Unify hiding (lookupVar)
import Text.Megaparsec (SourcePos)

-------------------------------
-- Lookup the value of TcEnv --
-------------------------------

lookupVar :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> RnId -> m (Scheme UType)
lookupVar pos name =
  use (varEnv . at name) >>= \case
    Nothing -> errorOn pos $ "Not in scope:" <+> quotes (pPrint name)
    Just scheme -> pure scheme

lookupType :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> RnId -> m UType
lookupType pos name =
  preuse (typeEnv . at name . _Just . typeConstructor) >>= \case
    Nothing -> errorOn pos $ "Not in scope:" <+> quotes (pPrint name)
    Just typ -> pure typ

lookupRecordType :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> [RnId] -> m (Scheme UType)
lookupRecordType pos fields = do
  env <- use fieldEnv
  case asumMap (`HashMap.lookup` env) fields of
    Nothing -> errorOn pos $ "Not in scope:" <+> (fields & map pPrint & punctuate " or" & sep)
    Just scheme -> pure scheme

typeCheck :: (MonadFail m, HasOpt env, MonadReader env m, MonadIO m, HasUniqSupply env, HasLogFunc env) => RnEnv -> Module (Malgo 'Rename) -> m (Module (Malgo 'TypeCheck), TcEnv)
typeCheck rnEnv (Module name bg) =
  runTypeUnifyT $ do
    tcEnv <- genTcEnv rnEnv
    (bg', tcEnv') <- runStateT (tcBindGroup bg) tcEnv
    -- FIXME: 自由なUVarに適当なTyVarを束縛する
    -- Right x |> { Right x -> print_Int32 x } みたいなパターンで必要
    zonkedBg <-
      pure bg'
        >>= traverseOf (scDefs . traversed . traversed . _1 . ann) zonk
        >>= traverseOf (scDefs . traversed . traversed . _3) (walkOn @TypeF @TypeVar zonk)
        >>= traverseOf (foreigns . traversed . _1 . ann) zonk
    zonkedTcEnv <-
      pure tcEnv'
        >>= traverseOf (varEnv . traversed . traversed) (walkOn @TypeF @TypeVar zonk)
        >>= traverseOf (typeEnv . traversed . traversed) (walkOn @TypeF @TypeVar zonk)
    pure (Module name zonkedBg, zonkedTcEnv)

tcBindGroup ::
  ( MonadState TcEnv m,
    MonadBind UType m,
    MonadFail m,
    MonadIO m,
    HasOpt env,
    MonadReader env m,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  BindGroup (Malgo 'Rename) ->
  m (BindGroup (Malgo 'TypeCheck))
tcBindGroup bindGroup = do
  _imports <- tcImports $ bindGroup ^. imports
  (_typeSynonyms, _dataDefs) <- tcTypeDefinitions (bindGroup ^. typeSynonyms) (bindGroup ^. dataDefs)
  _foreigns <- tcForeigns $ bindGroup ^. foreigns
  _scSigs <- tcScSigs $ bindGroup ^. scSigs
  traverse_ prepareTcScDefs $ bindGroup ^. scDefs
  _scDefs <- tcScDefGroup $ bindGroup ^. scDefs
  pure BindGroup {..}

tcImports ::
  ( MonadState TcEnv m,
    MonadIO m,
    HasOpt env,
    MonadReader env m,
    HasLogFunc env
  ) =>
  [Import (Malgo 'Rename)] ->
  m [Import (Malgo 'TypeCheck)]
tcImports = traverse tcImport
  where
    tcImport (pos, modName) = do
      interface <-
        loadInterface modName >>= \case
          Just x -> pure x
          Nothing -> errorOn pos $ "module" <+> pPrint modName <+> "is not found"
      varEnv <>= fmap (fmap Static.fromType) (interface ^. signatureMap)
      typeEnv <>= fmap (fmap Static.fromType) (interface ^. typeDefMap)
      pure (pos, modName)

tcTypeDefinitions ::
  ( MonadBind UType m,
    MonadState TcEnv m,
    MonadReader env m,
    MonadIO m,
    HasOpt env,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  [TypeSynonym (Malgo 'Rename)] ->
  [DataDef (Malgo 'Rename)] ->
  m ([TypeSynonym (Malgo 'TypeCheck)], [DataDef (Malgo 'TypeCheck)])
tcTypeDefinitions typeSynonyms dataDefs = do
  -- 相互再帰的な型定義がありうるため、型コンストラクタに対応するTyConを先にすべて生成する
  for_ typeSynonyms \(x, name, params, _) -> do
    tyCon <- UVar <$> freshVar @UType
    tyConKind <- kindOf tyCon
    solve [With x $ tyConKind :~ buildTyConKind params]
    typeEnv . at name .= Just (TypeDef tyCon [] [])
  for_ dataDefs \(_, name, params, _) -> do
    tyCon <- TyCon <$> newIdOnName (buildTyConKind params) name
    typeEnv . at name .= Just (TypeDef tyCon [] [])
  (,) <$> tcTypeSynonyms typeSynonyms
    <*> tcDataDefs dataDefs
  where
    -- TODO: ほんとはpolymorphicな値を返さないといけないと思う
    buildTyConKind [] = TYPE $ Rep BoxedRep
    buildTyConKind (_ : xs) = TyArr (TYPE $ Rep BoxedRep) (buildTyConKind xs)

tcTypeSynonyms ::
  ( MonadBind UType f,
    MonadState TcEnv f,
    MonadIO f,
    MonadReader env f,
    HasOpt env,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  [TypeSynonym (Malgo 'Rename)] ->
  f [TypeSynonym (Malgo 'TypeCheck)]
tcTypeSynonyms ds =
  for ds \(pos, name, params, typ) -> do
    logDebug $ displayShow $ "Parametized type synonym:" <+> pPrint (params, typ)
    -- unless (null params) do
    --   errorOn pos $
    --     "Parametized type synonym is not supported"
    --       $+$ "TODO: add type operator and fix TyTyple and TyLazy's kinding"
    name' <- lookupType pos name
    params' <- traverse (const $ UVar <$> freshVar @UType) params
    nameKind <- kindOf name'
    paramKinds <- traverse kindOf params'
    logDebug $ displayShow $ "unify" <+> parens (pPrint (foldr TyArr (TYPE $ Rep BoxedRep) paramKinds)) <+> parens (pPrint nameKind)
    solve [With pos $ foldr TyArr (TYPE $ Rep BoxedRep) paramKinds :~ nameKind]
    zipWithM_ (\p p' -> typeEnv . at p .= Just (TypeDef p' [] [])) params params'
    transedTyp <- transType typ
    logDebug $ displayShow $ "unify" <+> parens (pPrint $ buildTyApp name' params') <+> parens (pPrint transedTyp)
    solve [With pos $ buildTyApp name' params' :~ transedTyp]
    updateFieldEnv (tcType typ) [] transedTyp<commit_msg>Revert comment out<commit_after>module Malgo.TypeCheck.Pass where

import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import Data.List (find, head, last, splitAt)
import Data.List.Extra (anySame)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromJust)
import Koriel.Id
import Koriel.MonadUniq
import Koriel.Pretty
import Malgo.Interface (loadInterface, signatureMap, typeDefMap)
import Malgo.Prelude
import Malgo.Rename.RnEnv (RnEnv)
import qualified Malgo.Rename.RnEnv as R
import Malgo.Syntax hiding (Type (..), freevars)
import qualified Malgo.Syntax as S
import Malgo.Syntax.Extension
import Malgo.TypeCheck.TcEnv
import Malgo.TypeRep.Static (Rep (..), Scheme (Forall), TypeDef (..), TypeF, typeConstructor, typeParameters, valueConstructors)
import qualified Malgo.TypeRep.Static as Static
import Malgo.TypeRep.UTerm
import Malgo.UTerm
import Malgo.Unify hiding (lookupVar)
import Text.Megaparsec (SourcePos)

-------------------------------
-- Lookup the value of TcEnv --
-------------------------------

lookupVar :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> RnId -> m (Scheme UType)
lookupVar pos name =
  use (varEnv . at name) >>= \case
    Nothing -> errorOn pos $ "Not in scope:" <+> quotes (pPrint name)
    Just scheme -> pure scheme

lookupType :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> RnId -> m UType
lookupType pos name =
  preuse (typeEnv . at name . _Just . typeConstructor) >>= \case
    Nothing -> errorOn pos $ "Not in scope:" <+> quotes (pPrint name)
    Just typ -> pure typ

lookupRecordType :: (MonadState TcEnv m, HasOpt env, MonadReader env m, MonadIO m) => SourcePos -> [RnId] -> m (Scheme UType)
lookupRecordType pos fields = do
  env <- use fieldEnv
  case asumMap (`HashMap.lookup` env) fields of
    Nothing -> errorOn pos $ "Not in scope:" <+> (fields & map pPrint & punctuate " or" & sep)
    Just scheme -> pure scheme

typeCheck :: (MonadFail m, HasOpt env, MonadReader env m, MonadIO m, HasUniqSupply env, HasLogFunc env) => RnEnv -> Module (Malgo 'Rename) -> m (Module (Malgo 'TypeCheck), TcEnv)
typeCheck rnEnv (Module name bg) =
  runTypeUnifyT $ do
    tcEnv <- genTcEnv rnEnv
    (bg', tcEnv') <- runStateT (tcBindGroup bg) tcEnv
    -- FIXME: 自由なUVarに適当なTyVarを束縛する
    -- Right x |> { Right x -> print_Int32 x } みたいなパターンで必要
    zonkedBg <-
      pure bg'
        >>= traverseOf (scDefs . traversed . traversed . _1 . ann) zonk
        >>= traverseOf (scDefs . traversed . traversed . _3) (walkOn @TypeF @TypeVar zonk)
        >>= traverseOf (foreigns . traversed . _1 . ann) zonk
    zonkedTcEnv <-
      pure tcEnv'
        >>= traverseOf (varEnv . traversed . traversed) (walkOn @TypeF @TypeVar zonk)
        >>= traverseOf (typeEnv . traversed . traversed) (walkOn @TypeF @TypeVar zonk)
    pure (Module name zonkedBg, zonkedTcEnv)

tcBindGroup ::
  ( MonadState TcEnv m,
    MonadBind UType m,
    MonadFail m,
    MonadIO m,
    HasOpt env,
    MonadReader env m,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  BindGroup (Malgo 'Rename) ->
  m (BindGroup (Malgo 'TypeCheck))
tcBindGroup bindGroup = do
  _imports <- tcImports $ bindGroup ^. imports
  (_typeSynonyms, _dataDefs) <- tcTypeDefinitions (bindGroup ^. typeSynonyms) (bindGroup ^. dataDefs)
  _foreigns <- tcForeigns $ bindGroup ^. foreigns
  _scSigs <- tcScSigs $ bindGroup ^. scSigs
  traverse_ prepareTcScDefs $ bindGroup ^. scDefs
  _scDefs <- tcScDefGroup $ bindGroup ^. scDefs
  pure BindGroup {..}

tcImports ::
  ( MonadState TcEnv m,
    MonadIO m,
    HasOpt env,
    MonadReader env m,
    HasLogFunc env
  ) =>
  [Import (Malgo 'Rename)] ->
  m [Import (Malgo 'TypeCheck)]
tcImports = traverse tcImport
  where
    tcImport (pos, modName) = do
      interface <-
        loadInterface modName >>= \case
          Just x -> pure x
          Nothing -> errorOn pos $ "module" <+> pPrint modName <+> "is not found"
      varEnv <>= fmap (fmap Static.fromType) (interface ^. signatureMap)
      typeEnv <>= fmap (fmap Static.fromType) (interface ^. typeDefMap)
      pure (pos, modName)

tcTypeDefinitions ::
  ( MonadBind UType m,
    MonadState TcEnv m,
    MonadReader env m,
    MonadIO m,
    HasOpt env,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  [TypeSynonym (Malgo 'Rename)] ->
  [DataDef (Malgo 'Rename)] ->
  m ([TypeSynonym (Malgo 'TypeCheck)], [DataDef (Malgo 'TypeCheck)])
tcTypeDefinitions typeSynonyms dataDefs = do
  -- 相互再帰的な型定義がありうるため、型コンストラクタに対応するTyConを先にすべて生成する
  for_ typeSynonyms \(x, name, params, _) -> do
    tyCon <- UVar <$> freshVar @UType
    tyConKind <- kindOf tyCon
    solve [With x $ tyConKind :~ buildTyConKind params]
    typeEnv . at name .= Just (TypeDef tyCon [] [])
  for_ dataDefs \(_, name, params, _) -> do
    tyCon <- TyCon <$> newIdOnName (buildTyConKind params) name
    typeEnv . at name .= Just (TypeDef tyCon [] [])
  (,) <$> tcTypeSynonyms typeSynonyms
    <*> tcDataDefs dataDefs
  where
    -- TODO: ほんとはpolymorphicな値を返さないといけないと思う
    buildTyConKind [] = TYPE $ Rep BoxedRep
    buildTyConKind (_ : xs) = TyArr (TYPE $ Rep BoxedRep) (buildTyConKind xs)

tcTypeSynonyms ::
  ( MonadBind UType f,
    MonadState TcEnv f,
    MonadIO f,
    MonadReader env f,
    HasOpt env,
    HasUniqSupply env,
    HasLogFunc env
  ) =>
  [TypeSynonym (Malgo 'Rename)] ->
  f [TypeSynonym (Malgo 'TypeCheck)]
tcTypeSynonyms ds =
  for ds \(pos, name, params, typ) -> do
    logDebug $ displayShow $ "Parametized type synonym:" <+> pPrint (params, typ)
    unless (null params) do
      errorOn pos $
        "Parametized type synonym is not supported"
          $+$ "TODO: add type operator and fix TyTyple and TyLazy's kinding"
    name' <- lookupType pos name
    params' <- traverse (const $ UVar <$> freshVar @UType) params
    nameKind <- kindOf name'
    paramKinds <- traverse kindOf params'
    logDebug $ displayShow $ "unify" <+> parens (pPrint (foldr TyArr (TYPE $ Rep BoxedRep) paramKinds)) <+> parens (pPrint nameKind)
    solve [With pos $ foldr TyArr (TYPE $ Rep BoxedRep) paramKinds :~ nameKind]
    zipWithM_ (\p p' -> typeEnv . at p .= Just (TypeDef p' [] [])) params params'
    transedTyp <- transType typ
    logDebug $ displayShow $ "unify" <+> parens (pPrint $ buildTyApp name' params') <+> parens (pPrint transedTyp)
    solve [With pos $ buildTyApp name' params' :~ transedTyp]
    updateFieldEnv (tcType typ) [] transedTyp
{'fraction_non_alphanumeric': 0.12142358688066993, 'fraction_numerical': 0.0007753741180119407, 'mean_word_length': 3.2054124551679166, 'pattern_counts': {'":': 0, '<': 75, '<?xml version=': 0, '>': 139, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6180835', 'n_tokens_mistral': 4644, 'n_tokens_neox': 4517, 'n_words': 1343}