Var.hs

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

-- | Backend statements for variable manipulation: assignment, replacement, update.

module Indigo.Backend.Var
  ( assignVar
  , setVar
  , setField
  , updateVar
  ) where

import Fmt (pretty)

import Indigo.Backend.Expr.Compilation (binaryOpFlat, compileExpr, unaryOpFlat)
import Indigo.Backend.Expr.Decompose (ExprDecomposition(Deconstructed, ExprFields), decomposeExpr)
import Indigo.Backend.Lookup (varActionSet, varActionUpdate)
import Indigo.Backend.Prelude
import Indigo.Backend.Scope
import Indigo.Common.Expr (Expr(V))
import Indigo.Common.Field (FieldLens(DeeperField, TargetField), HasField(..), fetchField, flSFO)
import Indigo.Common.Object
import Indigo.Common.State
import Indigo.Common.Var (RefId, Var(..))
import Indigo.Lorentz
import Lorentz.Instr qualified as L
import Morley.Michelson.Typed.Haskell.Instr.Product (GetFieldType)
import Morley.Util.Type (type (++))

-- | Assign the given variable to the value resulting from the given expression.
assignVar :: forall x inp . KnownValue x => Var x -> Expr x -> IndigoState inp (x : inp)
assignVar var e =
  stmtHookState (prettyAssign @(Var x) var (pretty e)) $
    compileExpr e >> assignTopVar var

-- | Set the variable to a new value.
--
-- If a variable is a cell on the stack,
-- we just compile passed expression and replace variable cell on stack.
-- If a variable is decomposed, we decompose passed expression
-- and call 'setVar' recursively from its fields.
--
-- Pay attention that this function takes a next RefId but it doesn't return RefId
-- because all allocated variables will be destroyed during execution of the function,
-- so allocated ones won't affect next allocated ones.
setVar :: forall a inp. KnownValue a => RefId -> Var a -> Expr a -> IndigoState inp inp
setVar nextRef v ex = stmtHookState (prettyAssign @(Var a) v (pretty ex)) $
  withObjectState v $ flip (setVarImpl nextRef) ex

setVarImpl :: forall a inp . RefId -> Object a -> Expr a -> IndigoState inp inp
setVarImpl _ (Cell refId) ex = IndigoState $ \md -> usingIndigoState md $
  unaryOpFlat ex $ varActionSet refId (mdStack md)
setVarImpl nextRef (Decomposed fields) ex = IndigoState $ \md@MetaData{..} ->
  case decomposeExpr mdObjects ex of
    ExprFields fieldsExpr -> usingIndigoState md $
      rmapZipM (namedToTypedRec @a namedToTypedFieldObj fields) fieldsExpr
    Deconstructed comp ->
      let GenCode decomposeSt decomposeExCd _ = usingIndigoState md comp
          setAllFieldsCd =
            setFieldsOnStack
              (MetaData decomposeSt mdObjects mdHooks)
              (namedToTypedRec @a namedToTypedFieldObj fields) in
      GenCode mdStack (decomposeExCd # setAllFieldsCd) L.nop
  where
    -- Set fields, if they are decomposed on stack.
    setFieldsOnStack
      :: forall rs .
         MetaData (rs ++ inp)
      -> Rec TypedFieldObj rs
      -> (rs ++ inp) :-> inp
    setFieldsOnStack _ RNil = L.nop
    setFieldsOnStack md (TypedFieldObj f :& vs) =
      let tmpFieldVar = Var nextRef
          setVarMd = pushRefMd tmpFieldVar (popNoRefMd md) in
      (gcCode $ usingIndigoState setVarMd $ setVarImpl (nextRef + 1) f (V tmpFieldVar)) #
      L.drop #
      setFieldsOnStack (popNoRefMd md) vs

    -- Take list of fields (variables, referring to them)
    -- and list of corresponding expressions and call 'setVarImpl' recursively.
    rmapZipM :: Rec TypedFieldObj rs -> Rec Expr rs -> IndigoState inp inp
    rmapZipM RNil RNil = nopState
    rmapZipM (TypedFieldObj f :& flds) (e :& exprs) =
      setVarImpl nextRef f e >>
      rmapZipM flds exprs

-- | Set the field (direct or indirect) of a complex object.
setField
  :: forall dt fname ftype inp .
     ( IsObject dt
     , IsObject ftype
     , HasField dt fname ftype
     )
  => RefId -> Var dt -> Label fname -> Expr ftype -> IndigoState inp inp
setField nextRef v targetLb e =
  stmtHookState ("setField " <> pretty v <> " #" <> pretty targetLb <> " " <> pretty e) $
    withObjectState v setFieldImpl
  where
    setFieldImpl :: forall x . (IsObject x, HasField x fname ftype) => Object x -> IndigoState inp inp
    setFieldImpl (Cell refId) =
      updateVar @x nextRef (sopSetField (flSFO fieldLens)
        (fieldNameFromLabel targetLb)) (Var refId) e
    setFieldImpl (Decomposed fields) = case fieldLens @x @fname @ftype of
      TargetField lb _ ->
        case fetchField @x lb fields of
          NamedFieldObj field ->
            setVarImpl nextRef field e
      DeeperField (lb :: Label fnameInterm) _ ->
        case fetchField @x lb fields of
          NamedFieldObj vf ->
            setFieldImpl @(GetFieldType x fnameInterm) vf

-- | Call binary operator with constant argument to update a variable in-place.
updateVar
  :: forall x y inp . (IsObject x, KnownValue y)
  => RefId
  -> [y, x] :-> '[x]
  -> Var x
  -> Expr y
  -> IndigoState inp inp
updateVar nextRef action vr e =
  stmtHookState ("updateVar " <> pretty vr <> " with expr " <> pretty e) $
    withObjectState vr updateVarImpl
  where
    updateVarImpl (Cell refId) = IndigoState $ \md ->
      usingIndigoState md $ unaryOpFlat e $ varActionUpdate refId (mdStack md) action
    -- This function doesn't have to be called for complex data types,
    -- it's only supposed to be used for assign-like statements
    -- (+=), (-=), etc but implemented just in case.
    updateVarImpl obj@(Decomposed _) = IndigoState $ \md ->
      let tmpVar = Var nextRef in
      let newMd = pushRefMd tmpVar md in
      usingIndigoState md $ binaryOpFlat e (V vr) $
        L.framed action #
        gcCode (usingIndigoState newMd (setVarImpl (nextRef + 1) obj (V tmpVar))) #
        L.drop