{-# OPTIONS_GHC -Wno-name-shadowing #-}

-- | Intermediate representation (IR) stages of the compiler pipeline.
module IR where

import Common.Compiler
import Common.Default (Default (..))
import qualified Front.Ast as A
import qualified IR.IR as I
import qualified IR.Pretty ()

import Control.Monad (
  when,
  (>=>),
 )
import IR.ClassInstantiation (instProgram)
import IR.Constraint.Typechecking
import IR.DConToFunc (dConToFunc)
import IR.DesugarPattern (desugarPattern)
import IR.ExternToCall (externToCall)
import IR.InsertRefCounting (insertRefCounting)
import IR.LambdaLift (liftProgramLambdas)
import IR.LowerAst (lowerProgram)
import IR.MangleNames (mangleProgram)
import IR.OptimizePar (optimizePar)
import IR.Pattern (checkAnomaly)
import IR.SegmentLets (segmentLets)
-- import IR.Simplify (simplifyProgram)
import IR.Types (fromAnnotations)
import System.Console.GetOpt (
  ArgDescr (..),
  OptDescr (..),
 )
import Text.Show.Pretty


{- | Operation modes for the IR compiler stage.

 By default, 'Continue' completes the pipeline end-to-end.
-}
data Mode
  = Continue
  | DumpIR
  | DumpIRAnnotated
  | DumpIRConstraints
  | DumpIRTyped
  | DumpIRTypedUgly
  | DumpIRMangled
  | DumpIRInlined
  | DumpIRTypedShow
  | DumpIRLifted
  | DumpIRFinal
  deriving (Mode -> Mode -> Bool
(Mode -> Mode -> Bool) -> (Mode -> Mode -> Bool) -> Eq Mode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Mode -> Mode -> Bool
$c/= :: Mode -> Mode -> Bool
== :: Mode -> Mode -> Bool
$c== :: Mode -> Mode -> Bool
Eq, Int -> Mode -> ShowS
[Mode] -> ShowS
Mode -> String
(Int -> Mode -> ShowS)
-> (Mode -> String) -> ([Mode] -> ShowS) -> Show Mode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Mode] -> ShowS
$cshowList :: [Mode] -> ShowS
show :: Mode -> String
$cshow :: Mode -> String
showsPrec :: Int -> Mode -> ShowS
$cshowsPrec :: Int -> Mode -> ShowS
Show)


-- | Compiler options for the IR compiler stage.
newtype Options = Options {Options -> Mode
mode :: Mode}
  deriving (Options -> Options -> Bool
(Options -> Options -> Bool)
-> (Options -> Options -> Bool) -> Eq Options
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Options -> Options -> Bool
$c/= :: Options -> Options -> Bool
== :: Options -> Options -> Bool
$c== :: Options -> Options -> Bool
Eq, Int -> Options -> ShowS
[Options] -> ShowS
Options -> String
(Int -> Options -> ShowS)
-> (Options -> String) -> ([Options] -> ShowS) -> Show Options
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Options] -> ShowS
$cshowList :: [Options] -> ShowS
show :: Options -> String
$cshow :: Options -> String
showsPrec :: Int -> Options -> ShowS
$cshowsPrec :: Int -> Options -> ShowS
Show)


instance Default Options where
  def :: Options
def = Options :: Mode -> Options
Options{mode :: Mode
mode = Mode
Continue}


-- | CLI options for the IR compiler stage.
options :: [OptDescr (Options -> Options)]
options :: [OptDescr (Options -> Options)]
options =
  [ String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRAnnotated)
      String
"Print the IR immediately after lowering"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-annotated"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRTyped)
      String
"Print the fully-typed IR just before type inference"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-constraints"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRConstraints)
      String
"Print the constraint IR used by the constraint solver type inference"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-typed"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRTyped)
      String
"Print the fully-typed IR after type inference"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-typed-ugly"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRTypedShow)
      String
"Ugly-Print the fully-typed IR after type inference"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-mangled"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRMangled)
      String
"Print the IR after mangling"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-lifted"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRLifted)
      String
"Print the IR after lambda lifting"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-inlined"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRInlined)
      String
"Print IR after inlining optimization and before dup drops"
  , String
-> [String]
-> ArgDescr (Options -> Options)
-> String
-> OptDescr (Options -> Options)
forall a. String -> [String] -> ArgDescr a -> String -> OptDescr a
Option
      String
""
      [String
"dump-ir-final"]
      ((Options -> Options) -> ArgDescr (Options -> Options)
forall a. a -> ArgDescr a
NoArg ((Options -> Options) -> ArgDescr (Options -> Options))
-> (Options -> Options) -> ArgDescr (Options -> Options)
forall a b. (a -> b) -> a -> b
$ Mode -> Options -> Options
setMode Mode
DumpIRFinal)
      String
"Print the last IR representation before code generation"
  ]
 where
  setMode :: Mode -> Options -> Options
setMode Mode
m Options
o = Options
o{mode :: Mode
mode = Mode
m}


-- | Lower from AST to IR (with annotations).
lower :: Options -> A.Program -> Pass (I.Program I.Annotations)
lower :: Options -> Program -> Pass (Program Annotations)
lower Options
opt Program
p = do
  Program Annotations
p <- Program -> Pass (Program Annotations)
lowerProgram Program
p
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIR) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump (Program Type -> Pass ()) -> Program Type -> Pass ()
forall a b. (a -> b) -> a -> b
$ (Annotations -> Type) -> Program Annotations -> Program Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Annotations -> Type
fromAnnotations Program Annotations
p
  Program Annotations -> Pass (Program Annotations)
forall (m :: * -> *) a. Monad m => a -> m a
return Program Annotations
p


{- | Type inference + check against type annotations.

 After this stage, no compiler errors should be thrown.
-}
typecheck :: Options -> I.Program I.Annotations -> Pass (I.Program I.Type)
typecheck :: Options -> Program Annotations -> Pass (Program Type)
typecheck Options
opt Program Annotations
p = do
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRAnnotated) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump (Program Type -> Pass ()) -> Program Type -> Pass ()
forall a b. (a -> b) -> a -> b
$ (Annotations -> Type) -> Program Annotations -> Program Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Annotations -> Type
fromAnnotations Program Annotations
p
  Program Type
p <- Program Annotations -> Bool -> Pass (Program Type)
typecheckProgram Program Annotations
p (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRConstraints)
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRTyped) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump Program Type
p
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRTypedShow) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ (Error -> Pass ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Error -> Pass ())
-> (Program Type -> Error) -> Program Type -> Pass ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Error
Dump (String -> Error)
-> (Program Type -> String) -> Program Type -> Error
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Program Type -> String
forall a. Show a => a -> String
ppShow) Program Type
p
  Program Type -> Pass (Program Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Program Type
p


anomalycheck :: I.Program I.Type -> Pass (I.Program I.Type)
anomalycheck :: Program Type -> Pass (Program Type)
anomalycheck Program Type
p = do
  Program Type -> Pass ()
forall t. Show t => Program t -> Pass ()
checkAnomaly Program Type
p
  Program Type -> Pass (Program Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Program Type
p


-- | IR transformations to prepare for codegen.
transform :: Options -> I.Program I.Type -> Pass (I.Program I.Type)
transform :: Options -> Program Type -> Pass (Program Type)
transform Options
opt Program Type
p = do
  Program Type
p <- Program Type -> Pass (Program Type)
mangleProgram Program Type
p
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRMangled) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
desugarPattern Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
instProgram Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
segmentLets Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
dConToFunc Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
externToCall Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
optimizePar Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
liftProgramLambdas Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
segmentLets Program Type
p
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRLifted) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump Program Type
p
  -- p <- simplifyProgram p -- TODO: inline BEFORE lambda lifting!!
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRInlined) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump Program Type
p
  Program Type
p <- Program Type -> Pass (Program Type)
insertRefCounting Program Type
p
  Bool -> Pass () -> Pass ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Options -> Mode
mode Options
opt Mode -> Mode -> Bool
forall a. Eq a => a -> a -> Bool
== Mode
DumpIRFinal) (Pass () -> Pass ()) -> Pass () -> Pass ()
forall a b. (a -> b) -> a -> b
$ Program Type -> Pass ()
forall a x. Pretty a => a -> Pass x
dump Program Type
p
  Program Type -> Pass (Program Type)
forall (m :: * -> *) a. Monad m => a -> m a
return Program Type
p


-- | IR compiler stage.
run :: Options -> A.Program -> Pass (I.Program I.Type)
run :: Options -> Program -> Pass (Program Type)
run Options
opt = Options -> Program -> Pass (Program Annotations)
lower Options
opt (Program -> Pass (Program Annotations))
-> (Program Annotations -> Pass (Program Type))
-> Program
-> Pass (Program Type)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Options -> Program Annotations -> Pass (Program Type)
typecheck Options
opt (Program Annotations -> Pass (Program Type))
-> (Program Type -> Pass (Program Type))
-> Program Annotations
-> Pass (Program Type)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Program Type -> Pass (Program Type)
anomalycheck (Program Type -> Pass (Program Type))
-> (Program Type -> Pass (Program Type))
-> Program Type
-> Pass (Program Type)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Options -> Program Type -> Pass (Program Type)
transform Options
opt