-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{-# 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 qualified Data.Kind as Kind
import Fmt (Buildable(..), pretty)
import qualified GHC.TypeLits as Lit
import Util.Type (type (++))

import Indigo.Backend.Prelude
import Indigo.Internal.Expr hiding ((<>))
import Indigo.Internal.State
import Indigo.Internal.Var
import Indigo.Lorentz
import qualified Lorentz.Instr as L

-- | 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
  -- TODO is L.dip . L.dip cheaper than L.dipN ?
  liftClear' = L.dip . L.dip
  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