Compiler.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}

module Zwirn.Language.Compiler where

{-
    Compiler.hs - implementation of a compiler-interpreter for zwirn
    Copyright (C) 2023, Martin Gius

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this library.  If not, see <http://www.gnu.org/licenses/>.
-}

import Control.Concurrent (readMVar)
import Control.Concurrent.MVar (MVar, modifyMVar_)
import Control.Exception (SomeException, try)
import Control.Monad
import Control.Monad.Except
import Control.Monad.State
import Data.Either (lefts, rights)
import Data.List (intercalate)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map as Map
import Data.Text (Text, pack, unpack)
import qualified Data.Text as T
import Data.Text.IO (readFile)
import System.IO (hPutStrLn, stderr)
import Text.Read (readMaybe)
import Zwirn.Core.Types (Value (..), silence, toList, unzwirn, value)
import Zwirn.Language.Block
import Zwirn.Language.Builtin.Prelude (builtinEnvironmentWithPlayEnv, builtinNamesDoux, builtinNamesSuperDirt)
import Zwirn.Language.Environment
import Zwirn.Language.Evaluate
import Zwirn.Language.Location
import Zwirn.Language.Macro
import Zwirn.Language.Parser
import Zwirn.Language.Play (PlayEnv (..), renderStatus)
import Zwirn.Language.Pretty
import qualified Zwirn.Language.Rotate as R
import Zwirn.Language.Simple
import Zwirn.Language.Syntax
import Zwirn.Language.TypeCheck.Constraint
import Zwirn.Language.TypeCheck.Infer
import Zwirn.Language.TypeCheck.Types
import Prelude hiding (readFile)

newtype CIMessage
  = CIMessage Text
  deriving (Show, Eq)

data CurrentBlock
  = CurrentBlock Int Int
  deriving (Eq, Show)

data ConfigEnv
  = ConfigEnv
  { cConfigPath :: IO String,
    cResetConfig :: IO String
  }

data StreamType = SuperDirt | Doux deriving (Eq, Show)

data CiConfig = CiConfig
  { ciConfigOverwriteBuiltin :: Bool,
    ciConfigDynamicTypes :: Bool,
    ciConfigPrecision :: Rational,
    ciConfigStreamType :: StreamType
  }

data Environment
  = Environment
  { stateEnv :: MVar ExpressionMap,
    playEnv :: PlayEnv,
    intEnv :: InterpreterEnv,
    confEnv :: Maybe ConfigEnv,
    ciConfig :: CiConfig,
    macroMap :: MacroMap
  }

data ErrorType
  = ParseErr String RealSrcLoc
  | TypeErr TypeError
  | RotErr R.RotationError
  | OtherErr String
  | ManyErr [ErrorType]

instance Show ErrorType where
  show (ParseErr s _) = s
  show (TypeErr err) = unpack $ render err
  show (RotErr err) = show err
  show (OtherErr err) = err
  show (ManyErr errs) = intercalate "\n" $ map show errs

data CIError
  = CIError
  { eError :: ErrorType,
    eEnv :: Environment
  }

data CompilerOutput
  = OutMessage Text
  | OutEdits [CodeEdit]
  | NoOutput
  deriving (Eq, Show)

instance Show CIError where
  show (CIError err _) = show err

type CI = StateT Environment (ExceptT CIError IO)

runCI :: Environment -> CI a -> IO (Either CIError a)
runCI env m = runExceptT $ evalStateT m env

runCIEnv :: Environment -> CI a -> IO (Either CIError (a, Environment))
runCIEnv env m = runExceptT $ runStateT m env

debug :: (MonadIO m) => String -> m ()
debug msg = liftIO $ hPutStrLn stderr $ "[zwirnzi] " <> msg

compilerInterpreterBasic :: Text -> CI CompilerOutput
compilerInterpreterBasic input = do
  sy <- runParser input
  runSyntax True sy

compilerInterpreterWithBlock :: Int -> Text -> CI (CompilerOutput, Environment, (Int, Int))
compilerInterpreterWithBlock line input = do
  blocks <- runBlocks 0 input
  b <- runGetBlock line blocks
  sys <- parseBlock b
  r <- mapM (runSyntax True) sys
  e <- get
  return (last r, e, (getBlockStart b, getBlockEnd b))

getBlockStartEnd :: Int -> Text -> CI (Int, Int)
getBlockStartEnd line input = do
  blocks <- runBlocks 0 input
  b <- runGetBlock line blocks
  return (getBlockStart b, getBlockEnd b)

compilerInterpreterBlock :: Int -> Text -> CI (CompilerOutput, Environment)
compilerInterpreterBlock line input = compilerInterpreterWithBlock line input >>= \(co, env, _) -> return (co, env)

compilerInterpreterLine :: Int -> Text -> CI (CompilerOutput, Environment)
compilerInterpreterLine line input = do
  blocks <- runBlocks 0 input
  content <- runGetLine line blocks
  sy <- runParserWithPos line "" content
  r <- runSyntax True sy
  e <- get
  return (r, e)

compilerInterpreterWhole :: Text -> CI (CompilerOutput, Environment)
compilerInterpreterWhole input = do
  blocks <- runBlocks 0 input
  syss <- mapM parseBlock blocks
  rs <- mapM (runSyntax True) $ concat syss
  e <- get
  return (last rs, e)

compilerInterpreterBoot :: [Text] -> CI Environment
compilerInterpreterBoot ps = mapM_ (runSyntax False . Command . noLoc . LoadCommand) ps >> get

-----------------------------------------------------
----------------- Throwing Errors -------------------
-----------------------------------------------------

throw :: ErrorType -> CI a
throw err = do
  env <- get
  throwError $ CIError err env

-- | catches an error by wrapping the result back into Either
catch :: CI a -> CI (Either CIError a)
catch c = catchError (fmap Right c) (return . Left)

-- | sequences the given actions, accumulating any occuring errors and wrapping them into ManyErr
catchMany :: [CI a] -> CI [a]
catchMany cs = do
  es <- mapM catch cs
  let errs = lefts es
      vs = rights es
  case errs of
    [] -> return ()
    [e] -> throw $ eError e
    xs -> throw (ManyErr $ map eError xs)
  return vs

filterErrors :: [CI a] -> CI [a]
filterErrors cs = do
  es <- mapM catch cs
  return $ rights es

-----------------------------------------------------
---------------------- Parser -----------------------
-----------------------------------------------------

extractParseErrPos :: String -> ErrorType
extractParseErrPos s = case mpos of
  Just p -> ParseErr s p
  Nothing -> OtherErr s
  where
    (ls, rs) = break (== ',') $ drop 20 s
    (c1s, _) = break (== '\n') $ drop 9 rs
    mpos = do
      l <- readMaybe ls
      c1 <- readMaybe c1s
      return (RealSrcLoc "" l c1 l c1)

runParserWithPos :: Int -> Text -> Text -> CI Syntax
runParserWithPos ln srcp t = case parseSyntaxWithPos ln srcp t of
  Left err -> throw $ extractParseErrPos err
  Right s -> return s

runParser :: Text -> CI Syntax
runParser t = case parseSyntax t of
  Left err -> throw $ extractParseErrPos err
  Right s -> return s

runBlocks :: Int -> Text -> CI [Block]
runBlocks ln t = case parseBlocks ln t of
  Left err -> throw $ OtherErr err
  Right bs -> return bs

runGetBlock :: Int -> [Block] -> CI Block
runGetBlock i bs = case getBlock i bs of
  Left err -> throw $ OtherErr err
  Right b -> return b

runGetLine :: Int -> [Block] -> CI Text
runGetLine i bs = case getSingleLine i bs of
  Left err -> throw $ OtherErr err
  Right (Line _ _ c) -> return c

parseBlock :: Block -> CI [Syntax]
parseBlock (Block ls) = catchMany $ map (\l -> runParserWithPos (lStart l) "" (lContent l)) (NE.toList ls)

getSyntaxLine :: Int -> Text -> CI (Maybe Syntax)
getSyntaxLine ln input = do
  blocks <- runBlocks 0 input
  case getSingleLine ln blocks of
    Right line -> Just <$> runParserWithPos (lStart line) "" (lContent line)
    Left _ -> return Nothing

getSyntaxBlock :: Int -> Text -> CI [Syntax]
getSyntaxBlock ln input = do
  blocks <- runBlocks 0 input
  block <- runGetBlock ln blocks
  parseBlock block

-----------------------------------------------------
---------------------- Macro ------------------------
-----------------------------------------------------

macroCI :: LocTerm -> CI (LocTerm, [CodeEdit])
macroCI t = do
  mmap <- gets macroMap
  case runMacros t mmap of
    Just x -> return x
    _ -> throw $ OtherErr "Unknown macro!"

-----------------------------------------------------
---------------------- Desugar ----------------------
-----------------------------------------------------

runSimplify :: LocTerm -> CI LocSimpleTerm
runSimplify t = return $ simplifyLoc t

runSimplifyDef :: Located Definition -> CI (Located SimpleDef, [CodeEdit])
runSimplifyDef (Located p (Definition x vs t)) = do
  (t', es) <- macroCI t
  return (Located p (LetS x (noLoc $ simplify $ TLambda vs t')), es)

-----------------------------------------------------
------------------- AST Rotation --------------------
-----------------------------------------------------

runRotate :: LocSimpleTerm -> CI LocSimpleTerm
runRotate s = case R.runRotate s of
  Left err -> throw $ RotErr err
  Right t -> return t

-----------------------------------------------------
-------------------- Type Check ---------------------
-----------------------------------------------------

runTypeCheck :: LocSimpleTerm -> CI Scheme
runTypeCheck s = do
  Environment {intEnv = env} <- get
  case inferTerm env s of
    Left err -> throw (TypeErr err)
    Right t -> return t

-----------------------------------------------------
-------------------- Interpreter --------------------
-----------------------------------------------------

interpret :: LocSimpleTerm -> CI Expression
interpret input = do
  env <- gets intEnv
  return $ evaluate env input

-- if ctx is false, highlighting should be disabled
checkHighlight :: Bool -> Expression -> CI Expression
checkHighlight True x = return x
checkHighlight False x = return $ removePosExp x

-----------------------------------------------------
----------------- Checking Options  -----------------
-----------------------------------------------------

overwriteOk :: Text -> CI ()
overwriteOk name = do
  overwrite <- gets (ciConfigOverwriteBuiltin . ciConfig)
  strtyp <- gets (ciConfigStreamType . ciConfig)
  let builtins = if strtyp == Doux then builtinNamesDoux else builtinNamesSuperDirt
  when (not overwrite && name `elem` builtins) $ throw $ OtherErr "Cannot overwrite builtin function. Please use OverwriteBuiltin."

dynamicOk :: Text -> Scheme -> CI ()
dynamicOk name ty = do
  dynamic <- gets (ciConfigDynamicTypes . ciConfig)
  mayty <- gets (lookupType name . intEnv)
  case mayty of
    Just oldType ->
      when (not dynamic && not (unifiable (schemeToType oldType, schemeToType ty))) $ throw $ OtherErr "Cannot overwrite definition with new type. Please use DynamicTypes."
    Nothing -> return ()

-------------------------------------------------------
----------------- Interpreting Syntax -----------------
-------------------------------------------------------

runSyntax :: Bool -> Syntax -> CI CompilerOutput
runSyntax b (Exec t) = executeTerm b t
runSyntax _ (Command c) = runCommand (lValue c)
runSyntax b (Def d) = define b d
runSyntax b (DynDef d) = dynamicDefine b (lValue d)
runSyntax _ (MacroDef d) = macroDefine (lValue d)

executeTerm :: Bool -> LocTerm -> CI CompilerOutput
executeTerm ctx t = do
  (t', es) <- macroCI t
  s <- runSimplify t'
  rot <- runRotate s
  ty <- runTypeCheck rot
  ex <- interpret rot
  exCtx <- checkHighlight ctx ex
  case ty of
    Forall _ (Qual _ _ (TypeCon "Action")) -> do
      stMV <- gets stateEnv
      liftIO $ evalAction stMV (fromExp exCtx)
      return $ OutEdits es
    _ -> throw $ OtherErr "Can only execute actions!"

define :: Bool -> Located Definition -> CI CompilerOutput
define ctx d = do
  (Located _ (LetS x st), es) <- runSimplifyDef d
  rot <- runRotate st
  ty <- runTypeCheck rot
  ex <- interpret rot
  exCtx <- checkHighlight ctx ex
  overwriteOk x
  dynamicOk x ty
  modify (\env -> env {intEnv = extend (x, exCtx, ty) (intEnv env)})
  return $ OutEdits es

dynamicDefine :: Bool -> DynamicDefinition -> CI CompilerOutput
dynamicDefine ctx (DynamicDefinition x t) = do
  (t', es) <- macroCI t
  s <- runSimplify t'
  rot <- runRotate s
  ty <- runTypeCheck rot
  ex <- interpret rot
  exCtx <- checkHighlight ctx ex
  overwriteOk x
  dynamicOk x ty
  setExpression x ty exCtx
  return $ OutEdits es

macroDefine :: MacroDefinition -> CI CompilerOutput
macroDefine (MacroDefinition x t) = do
  (t', es) <- macroCI t
  s <- runSimplify t'
  rot <- runRotate s
  _ <- runTypeCheck rot
  modify (\env -> env {macroMap = Map.insert x t' (macroMap env)})
  return $ OutEdits es

setExpression :: Text -> Scheme -> Expression -> CI ()
setExpression x ty exCtx
  | isBasicType ty = do
      if checkDependency x ty
        then do
          let newEx
                | isNumberT ty = EZwirn $ getStateN (pure x)
                | isTextT ty = EZwirn $ getStateT (pure x)
                | isMapT ty = EZwirn $ getStateM (pure x)
                | otherwise = EZwirn silence
          modify (\env -> env {intEnv = extend (x, newEx, addDependency x ty) (intEnv env)})
          stMV <- gets stateEnv
          liftIO $ stateSet stMV x exCtx
        else throw $ OtherErr "Cyclic dependency detected!"
  | otherwise = throw $ OtherErr "Can only set basic types!"

---------------------------------------------------------
----------------- Interpreting Commands -----------------
---------------------------------------------------------

runCommand :: Command -> CI CompilerOutput
runCommand (ShowCommand t) = showCommand t
runCommand (TypeCommand t) = typeCommand t
runCommand (SetCommand f) = setCommand (unpack f)
runCommand (UnsetCommand _) = return NoOutput
runCommand (LoadCommand f) = NoOutput <$ loadCommand f
runCommand (InfoCommand f) = infoCommand f
runCommand ResetEnvCommand = resetEnvCommand
runCommand ResetConfigCommand = resetConfigCommand
runCommand ShowConfigPathCommand = showConfigPathCommand
runCommand StatusCommand = statusCommand
runCommand EnvCommand = envCommand

showCommand :: LocTerm -> CI CompilerOutput
showCommand t = do
  (t', _) <- macroCI t
  s <- runSimplify t'
  rot <- runRotate s
  ty <- runTypeCheck rot
  if isBasicType ty
    then do
      ex <- interpret rot
      stMV <- gets stateEnv
      prec <- gets (ciConfigPrecision . ciConfig)
      st <- liftIO $ readMVar stMV
      return $ OutMessage $ pack $ showWithStatePrec (realToFrac prec) st ex
    else throw $ OtherErr $ "Can not show expressions of type: " ++ unpack (ppscheme ty)

typeCommand :: LocTerm -> CI CompilerOutput
typeCommand t = do
  (t', _) <- macroCI t
  s <- runSimplify t'
  rot <- runRotate s
  ty <- runTypeCheck rot
  return $ OutMessage $ ppTermHasType (t, ty)

loadCommand :: Text -> CI ()
loadCommand path = do
  mayfile <- liftIO ((try $ readFile $ unpack path) :: IO (Either SomeException Text))
  case mayfile of
    Left _ -> throw $ OtherErr "File not found"
    Right input -> do
      blocks <- runBlocks 0 input
      let content = concatMap getBlockContent blocks
      ss <- mapM runParser content
      mapM_ (runSyntax False) ss

infoCommand :: Text -> CI CompilerOutput
infoCommand n = do
  env <- gets intEnv
  case lookupFull n env of
    Just (Annotated _ t (Just d)) -> return $ OutMessage $ n <> " :: " <> ppscheme t <> "  \n" <> d
    Just (Annotated _ t Nothing) -> return $ OutMessage $ n <> " :: " <> ppscheme t
    Nothing -> return $ OutMessage $ pack $ "Couldn't find information about " ++ unpack n

resetConfigCommand :: CI CompilerOutput
resetConfigCommand = do
  (Environment {confEnv = mayEnv}) <- get
  case mayEnv of
    Nothing -> throw $ OtherErr "Configuration not available."
    Just (ConfigEnv _ reset) -> OutMessage . pack <$> liftIO reset

showConfigPathCommand :: CI CompilerOutput
showConfigPathCommand = do
  (Environment {confEnv = mayEnv}) <- get
  case mayEnv of
    Nothing -> throw $ OtherErr "Configuration not available."
    Just (ConfigEnv path _) -> OutMessage . pack <$> liftIO path

resetEnvCommand :: CI CompilerOutput
resetEnvCommand = do
  penv <- gets playEnv
  modify (\env -> env {intEnv = builtinEnvironmentWithPlayEnv penv})
  return $ OutMessage "Environment reset to default!"

setCommand :: String -> CI CompilerOutput
setCommand "DynamicTypes" = modify (\env -> env {ciConfig = (ciConfig env) {ciConfigDynamicTypes = True}}) >> return (OutMessage "Successfully enabled DynamicTypes.")
setCommand "OverwriteBuiltin" = modify (\env -> env {ciConfig = (ciConfig env) {ciConfigOverwriteBuiltin = True}}) >> return (OutMessage "Successfully enabled OverwriteBuiltin.")
setCommand _ = return $ OutMessage "Unknown compiler flag. The flags are: DynamicTypes, OverwriteBuiltin."

-- for now only prints out the "basic" non-builtin expressions
-- TODO: add a flag/modifier to print all expressions (via :env all)
envCommand :: CI CompilerOutput
envCommand = do
  env <- gets (Map.toList . Map.filter isBasicExpression . eExpressions . intEnv)
  penv <- gets playEnv
  let builtin = Map.keys $ Map.filter isBasicExpression $ eExpressions $ builtinEnvironmentWithPlayEnv penv
      filtered = filter (\(k, _) -> k `notElem` builtin) env
  return $ OutMessage $ T.intercalate "\n" $ map (\(k, Annotated _ ty _) -> k <> " :: " <> ppscheme ty) filtered

statusCommand :: CI CompilerOutput
statusCommand = do
  stMV <- gets stateEnv
  pmMV <- gets (playMap . playEnv)
  pm <- liftIO $ readMVar pmMV
  st <- liftIO $ readMVar stMV
  return $ OutMessage $ renderStatus (maybe 138 (getTempo st) (Map.lookup "_tempo" st)) pm
  where
    getTempo :: ExpressionMap -> Expression -> Double
    getTempo st x = case vs of
      [] -> 138
      ((Value v _ _, _) : _) -> v
      where
        z = fromExp x :: Zwirn Double
        vs = toList $ unzwirn z 0 st

isNumberT :: Scheme -> Bool
isNumberT (Forall _ (Qual _ _ (TypeCon "Number"))) = True
isNumberT _ = False

isTextT :: Scheme -> Bool
isTextT (Forall _ (Qual _ _ (TypeCon "Text"))) = True
isTextT _ = False

isMapT :: Scheme -> Bool
isMapT (Forall _ (Qual _ _ (TypeCon "Map"))) = True
isMapT _ = False

stateSet :: MVar ExpressionMap -> T.Text -> Expression -> IO ()
stateSet stMV x ex = modifyMVar_ stMV (return . Map.insert x ex)

updateState :: MVar ExpressionMap -> [ExpressionMap] -> IO ()
updateState _ [] = return ()
updateState stmv (st : _) = modifyMVar_ stmv (const $ return st)

evalAction :: MVar ExpressionMap -> Zwirn Expression -> IO ()
evalAction stMV z = do
  st <- readMVar stMV
  let exps = toList $ unzwirn z 0 st
      sts = map snd exps
      exs = map (value . fst) exps

  updateState stMV sts
  mapM_ evalActionExp exs
  where
    evalActionExp (EAction i) = i
    evalActionExp _ = return ()