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

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

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

import Fmt (pretty)

import Indigo.Backend.Prelude
import Indigo.Backend.Scope
import Indigo.Internal hiding ((+), (<>))
import Indigo.Lorentz
import qualified Lorentz.Instr as L
import Michelson.Typed.Haskell.Instr.Product (GetFieldType)
import 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) 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