Scope.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- | Machinery that provides the ability to return values from Indigo statements
-- (like @if@, @case@, @while@, etc).
-- You are allowed to return unit, one expression or a tuple of expressions.
-- For instance:
--
-- @
-- (a, b) <- if flag
--           then do
--                  anotherFlag <- newVar True
--                  return (5 +. var, anotherFlag ||. True)
--           else return (0, anotherVar)
-- @
-- is a valid construction.
-- Pay attention to the fact that @5 +. var@ has the type 'Expr' 'Integer',
-- but 0 is just an 'Integer' and @anotherFlag ||. True@ has type 'Expr' 'Bool',
-- but @anotherVar@ has type 'Var' 'Bool'; and this code will compile anyway.
-- This is done intentionally to avoid the burden of manually converting values
-- to expressions (or variables).
-- So you can write the same constructions as in a regular language.

module Indigo.Backend.Scope
  ( BranchRetKind (..)
  , ScopeCodeGen
  , ScopeCodeGen' (..)
  , ReturnableValue
  , ReturnableValue' (..)
  , RetOutStack
  , RetVars
  , RetExprs
  , ClassifyReturnValue
  , liftClear
  , compileScope
  , allocateVars
  , finalizeStatement

  -- Builder helpers for hooks
  , prettyAssign
  , condStmtPretty
  , prettyRet
  ) where

import Data.Kind qualified as Kind
import Fmt (Buildable(..), pretty)
import GHC.TypeLits qualified as Lit

import Indigo.Backend.Expr.Compilation (compileExpr)
import Indigo.Backend.Prelude
import Indigo.Common.Expr (Expr, ExprType, ToExpr, toExpr)
import Indigo.Common.State
import Indigo.Common.Var
import Indigo.Lorentz
import Lorentz.Instr qualified as L
import Morley.Util.Type (type (++))

-- | To avoid overlapping instances we need to somehow distinguish single values
-- from tuples, because the instances:
--
-- @
--   instance Something a
--   instance Something (a, b)
-- @
-- overlap and adding @&#x7b;-\# OVERLAPPING \#-&#x7d;@ doesn't rescue in some cases,
-- especially for type families defined in @Something@.
data BranchRetKind =
    Unit
  -- ^ If value is unit (don't return anything)
  | SingleVal
  -- ^ If it's a single value (not tuple)
  | Tuple
  -- ^ If it's tuple (we don't care how many elements are in)

-- | This type family returns a promoted value of type 'BranchRetKind'
-- or causes a compilation error if a tuple with too many elements is used.
type family ClassifyReturnValue (ret :: Kind.Type) where
  ClassifyReturnValue ()     = 'Unit
  ClassifyReturnValue (_, _) = 'Tuple
  -- These type errors are an attempt to make compilation errors clear
  -- in cases where one tries to return a tuple with more elements from a statement
  ClassifyReturnValue (_, _, _) = 'Tuple
  ClassifyReturnValue (_, _, _, _) =
    Lit.TypeError ('Lit.Text "Tuple with 4 elements is not supported yet as returning value")
  ClassifyReturnValue (_, _, _, _, _) =
    Lit.TypeError ('Lit.Text "Tuple with 5 elements is not supported yet as returning value")
  ClassifyReturnValue (_, _, _, _, _, _) =
    Lit.TypeError ('Lit.Text "Tuple with 6 elements is not supported yet as returning value")
  -- I hope nobody will try to return as a value tuples with more elements
  ClassifyReturnValue _      = 'SingleVal

-- | Class for values that can be returned from Indigo statements.
-- They include @()@ and tuples.
class ReturnableValue' (retKind :: BranchRetKind) (ret :: Kind.Type) where
  -- | Type family reflecting the top elements of stack produced by
  -- a statement returning the value.
  type family RetOutStack' retKind ret :: [Kind.Type]

  -- | Type family reflecting the returning value from a statement.
  type family RetVars' retKind ret :: Kind.Type

  -- | Tuple looking like @(Expr x, Expr y, ..)@ that corresponds
  -- to expressions returning from the scope.
  -- 'RetVars\'' and 'RetExprs\'' are twin types because
  -- the former just adds 'Var' over each expression of the latter.
  type family RetExprs' retKind ret :: Kind.Type

  -- | Allocate variables referring to result of the statement.
  -- Requires an allocator operating in a Monad.
  allocateVars'
    :: Monad m
    => (forall (x :: Kind.Type) . m (Var x))
    -> m (RetVars' retKind ret)

  -- | Push the variables referring to the result of the statement on top of
  -- the stack of the given 'StackVars'.
  assignVars'
    :: RetVars' retKind ret
    -> StackVars inp
    -> StackVars (RetOutStack' retKind ret ++ inp)

  -- | Pretty printing of statements like \"var := statement\"
  prettyAssign' :: RetVars' retKind ret -> Text -> Text

  -- | Prettify 'ret' value
  prettyRet' :: ret -> Text

-- | Type class which unions all related management of computations in a scope,
-- like in @if@ branch, in @case@ body, etc.
--
-- Particularly, it takes care of the computation of expressions returning
-- from a scope to leave it safely.
-- Basically, this type class encapsulates the generation of Lorentz code that looks like:
--
-- @
--   branch_code #
--     -- we get some arbitrary type of a stack here, lets call it @xs@
--   compute_returning_expressions #
--     -- we get type of stack [e1, e2, ... ek] ++ xs
--   cleanup_xs_to_inp
--     -- we get [e1, e2, e3, ..., ek] ++ inp
-- @
class ReturnableValue' retKind ret => ScopeCodeGen' (retKind :: BranchRetKind) (ret :: Kind.Type) where
  -- | Produces an Indigo computation that puts on the stack
  -- the evaluated returned expressions from the leaving scope.
  compileScopeReturn' :: ret -> IndigoState xs (RetOutStack' retKind ret ++ xs)

  -- | Drop the stack cells that were produced in the leaving scope,
  -- apart from ones corresponding to the returning expressions.
  liftClear' :: (xs :-> inp) -> (RetOutStack' retKind ret ++ xs :-> RetOutStack' retKind ret ++ inp)

  -- | Generate 'gcClear' for the whole statement
  genGcClear' :: (RetOutStack' retKind ret ++ inp) :-> inp

type RetOutStack ret = RetOutStack' (ClassifyReturnValue ret) ret
type RetVars ret = RetVars' (ClassifyReturnValue ret) ret
type RetExprs ret = RetExprs' (ClassifyReturnValue ret) ret
type ReturnableValue ret = ReturnableValue' (ClassifyReturnValue ret) ret
type ScopeCodeGen ret = ScopeCodeGen' (ClassifyReturnValue ret) ret

-- | Specific version of 'allocateVars\''
allocateVars
  :: forall ret m . (ReturnableValue ret, Monad m)
  => (forall (x :: Kind.Type) . m (Var x))
  -> m (RetVars ret)
allocateVars = allocateVars' @(ClassifyReturnValue ret) @ret

-- | Specific version of 'liftClear\''
liftClear
  :: forall ret inp xs . ScopeCodeGen ret
  => (xs :-> inp)
  -> (RetOutStack ret ++ xs :-> RetOutStack ret ++ inp)
liftClear = liftClear' @(ClassifyReturnValue ret) @ret

prettyAssign :: forall ret . ReturnableValue ret => RetVars ret -> Text -> Text
prettyAssign = prettyAssign' @(ClassifyReturnValue ret) @ret

prettyRet :: forall ret . ReturnableValue ret => ret -> Text
prettyRet = prettyRet' @(ClassifyReturnValue ret) @ret

condStmtPretty :: forall ret x . ReturnableValue ret => RetVars ret -> Text -> Expr x -> Text
condStmtPretty retVars stmtName ex = prettyAssign @ret retVars (stmtName <> " (" <> pretty ex <> ")")

-- | Concatenate a scoped code, generation of returning expressions,
-- and clean up of redundant cells from the stack.
compileScope
  :: forall ret inp xs . ScopeCodeGen ret
  => (StackVars xs -> MetaData xs)
  -- ^ Partially applied constructor of 'MetaData' (without passed 'StackVars').
  -- 'compileScope' function is usually being called from another function
  -- which is in 'IndigoState' and, consequently, holding 'MetaData' with all fields.
  -> GenCode inp xs
  -- ^ Code (and clear) of a wrapping scope
  -> ret
  -- ^ Return value of a scope (either primitives or expressions or variables)
  -> (inp :-> RetOutStack ret ++ inp)
compileScope mdCr innerGc gcRet =
  let md = mdCr (gcStack innerGc) in
  gcCode innerGc #
  auxiliaryHook md ("computation of returning values: " <> prettyRet gcRet)
    (gcCode $ usingIndigoState md $ compileScopeReturn' @(ClassifyReturnValue ret) gcRet) #
  auxiliaryHook md "dropping cells from the stack allocated in the scope"
    (liftClear' @(ClassifyReturnValue ret) @ret (gcClear innerGc))

-- | Push variables in the 'StackVars', referring to the generated expressions,
-- and generate 'gcClear' for the whole statement.
finalizeStatement
  :: forall ret inp . ScopeCodeGen ret
  => StackVars inp
  -> RetVars ret
  -> (inp :-> RetOutStack ret ++ inp)
  -> GenCode inp (RetOutStack ret ++ inp)
finalizeStatement md vars code =
  let newMd = assignVars' @(ClassifyReturnValue ret) @ret vars md in
  GenCode newMd code (genGcClear' @(ClassifyReturnValue ret) @ret)

-- Type instances for ScopeCodeGen'.
-- Perhaps, they could be implemented more succinctly
-- and expressed inductively via previous instances,
-- but I don't think it makes sense to spend a lot of time to shorten them.

type KnownValueExpr a = (KnownValue (ExprType a), ToExpr a)

instance ReturnableValue' 'Unit () where
  type RetOutStack' 'Unit () = '[]
  type RetVars' 'Unit () = ()
  type RetExprs' 'Unit () = ()
  allocateVars' _ = pure ()
  assignVars' _ md = md
  prettyAssign' _ stmt = stmt
  prettyRet' _ = "()"

instance ScopeCodeGen' 'Unit () where
  compileScopeReturn' _ = nopState
  liftClear' = id
  genGcClear' = L.nop

instance KnownValueExpr single => ReturnableValue' 'SingleVal single where
  type RetOutStack' 'SingleVal single = '[ExprType single]
  type RetVars' 'SingleVal single = Var (ExprType single)
  type RetExprs' 'SingleVal single = ExprType single
  allocateVars' allocator = allocator @(ExprType single)
  assignVars' = pushRef
  prettyAssign' retVars stmt = pretty retVars <> " := " <> stmt
  prettyRet' = pretty . toExpr

instance KnownValueExpr single => ScopeCodeGen' 'SingleVal single where
  compileScopeReturn' = compileToExpr
  liftClear' = L.dip
  genGcClear' = L.drop

instance ( KnownValueExpr x
         , KnownValueExpr y
         , Buildable (RetVars' 'Tuple (x, y))
         )
         => ReturnableValue' 'Tuple (x, y) where
  type RetOutStack' 'Tuple (x, y) = ExprType x ': '[ExprType y]
  type RetVars' 'Tuple (x, y) = (Var (ExprType x), Var (ExprType y))
  type RetExprs' 'Tuple (x, y) = (ExprType x, ExprType y)
  allocateVars' allocator = (,) <$> allocator <*> allocator
  assignVars' (var1, var2) md = pushRef var1 $ pushRef var2 md
  prettyAssign' retVars stmt = pretty retVars <> " := " <> stmt
  prettyRet' (x, y) = "(" <> pretty (toExpr x) <> ", " <> pretty (toExpr y) <> ")"

instance (KnownValueExpr x
         , KnownValueExpr y
         , Buildable (RetVars' 'Tuple (x, y))
         ) => ScopeCodeGen' 'Tuple (x, y) where
  compileScopeReturn' (e1, e2) = compileToExpr e2 >> compileToExpr e1
  liftClear' = L.dipN @2
  genGcClear' = L.drop # L.drop

instance ( KnownValueExpr x
         , KnownValueExpr y
         , KnownValueExpr z
         , Buildable (RetVars' 'Tuple (x, y, z))
         ) => ReturnableValue' 'Tuple (x, y, z) where
  type RetOutStack' 'Tuple (x, y, z) = ExprType x ': ExprType y ': '[ExprType z]
  type RetVars' 'Tuple (x, y, z) = (Var (ExprType x), Var (ExprType y), Var (ExprType z))
  type RetExprs' 'Tuple (x, y, z) = (ExprType x, ExprType y, ExprType z)
  allocateVars' allocator = (,,) <$> allocator <*> allocator <*> allocator
  assignVars' (var1, var2, var3) md =
    pushRef var1 . pushRef var2 $ pushRef var3 md
  prettyAssign' retVars stmt = pretty retVars <> " := " <> stmt
  prettyRet' (x, y, z) = "(" <> pretty (toExpr x) <> ", " <> pretty (toExpr y) <> ", " <> pretty (toExpr z) <> ")"

instance (KnownValueExpr x
         , KnownValueExpr y
         , KnownValueExpr z
         , Buildable (RetVars' 'Tuple (x, y, z))
         ) => ScopeCodeGen' 'Tuple (x, y, z) where
  compileScopeReturn' (e1, e2, e3) = compileToExpr e3 >> compileToExpr e2 >> compileToExpr e1
  liftClear' = L.dipN @3
  genGcClear' = L.drop # L.drop # L.drop

-- | Utility function to compile from an 'IsExpr'
compileToExpr :: ToExpr a => a -> IndigoState inp ((ExprType a) : inp)
compileToExpr = compileExpr . toExpr