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

-- | 'Expr' compilation

module Indigo.Internal.Expr.Compilation
  ( compileExpr

  , ObjManipulationRes (..)
  , runObjectManipulation
  , namedToExpr

  , nullaryOp
  , unaryOp
  , binaryOp
  , ternaryOp

  , nullaryOpFlat
  , unaryOpFlat
  , binaryOpFlat
  , ternaryOpFlat
  ) where

import Data.Vinyl.Core (RMap(..))

import qualified Lorentz.ADT as L
import qualified Lorentz.Instr as L
import qualified Lorentz.Macro as L
import qualified Lorentz.StoreClass as L
import Michelson.Typed.Haskell.Instr.Product (GetFieldType)

import Indigo.Backend.Prelude
import Indigo.Internal.Expr.Types
import Indigo.Internal.Field
import Indigo.Internal.Lookup (varActionGet)
import Indigo.Internal.Object
  (IndigoObjectF(..), NamedFieldObj(..), castFieldConstructors, namedToTypedRec, typedToNamedRec)
import Indigo.Internal.State
  (DecomposedObjects, GenCode(..), IndigoState(..), MetaData(..), replStkMd, usingIndigoState,
  withObject, withObjectState)
import Indigo.Internal.Var (Var(..), pushNoRef)
import Indigo.Lorentz

compileExpr :: forall a inp . Expr a -> IndigoState inp (a : inp)
compileExpr (C a) = IndigoState $ \md -> GenCode (pushNoRef $ mdStack md) (L.push a) L.drop
compileExpr (V v) = withObjectState v $ compileObjectF namedToExpr
compileExpr (Update m key val) = ternaryOp key val m L.update
compileExpr (Add e1 e2) = binaryOp e1 e2 L.add
compileExpr (Sub e1 e2) = binaryOp e1 e2 L.sub
compileExpr (Mul e1 e2) = binaryOp e1 e2 L.mul
compileExpr (Div e1 e2) = binaryOp e1 e2 (L.ediv # L.ifSome L.car (failUsing [mt|devision by zero|]))
compileExpr (Mod e1 e2) = binaryOp e1 e2 (L.ediv # L.ifSome L.cdr (failUsing [mt|devision by zero|]))
compileExpr (Abs e) = unaryOp e L.abs
compileExpr (Neg e) = unaryOp e L.neg

compileExpr (Lsl e1 e2) = binaryOp e1 e2 L.lsl
compileExpr (Lsr e1 e2) = binaryOp e1 e2 L.lsr

compileExpr (Eq' e1 e2) = binaryOp e1 e2 L.eq
compileExpr (Neq e1 e2) = binaryOp e1 e2 L.neq
compileExpr (Lt e1 e2) = binaryOp e1 e2 L.lt
compileExpr (Le e1 e2) = binaryOp e1 e2 L.le
compileExpr (Gt e1 e2) = binaryOp e1 e2 L.gt
compileExpr (Ge e1 e2) = binaryOp e1 e2 L.ge
compileExpr (IsNat e) = unaryOp e L.isNat
compileExpr (Int' e) = unaryOp e L.int
compileExpr (Coerce e) = unaryOp e checkedCoerce_
compileExpr (ForcedCoerce e) = unaryOp e forcedCoerce_
compileExpr (And e1 e2) = binaryOp e1 e2 L.and
compileExpr (Or e1 e2) = binaryOp e1 e2 L.or
compileExpr (Xor e1 e2) = binaryOp e1 e2 L.xor
compileExpr (Not e) = unaryOp e L.not

compileExpr (Fst e) = unaryOp e L.car
compileExpr (Snd e) = unaryOp e L.cdr
compileExpr (Pair e1 e2) = binaryOp e1 e2 L.pair

compileExpr (Some e) = unaryOp e L.some
compileExpr None = nullaryOp L.none
compileExpr (Right' e) = unaryOp e L.right
compileExpr (Left' e) = unaryOp e L.left
compileExpr (Pack e) = unaryOp e L.pack
compileExpr (Unpack e) = unaryOp e L.unpack
compileExpr (PackRaw e) = unaryOp e L.packRaw
compileExpr (UnpackRaw e) = unaryOp e L.unpackRaw
compileExpr Nil = nullaryOp L.nil
compileExpr (Cons e1 e2) = binaryOp e1 e2 L.cons
compileExpr (Contract e) = unaryOp e L.contract
compileExpr Self = nullaryOp L.self
compileExpr (ContractAddress ec) = unaryOp ec L.address
compileExpr (ContractCallingUnsafe epName addr) = unaryOp addr (L.contractCallingUnsafe epName)
compileExpr (RunFutureContract con) = unaryOp con L.runFutureContract
compileExpr (ConvertEpAddressToContract epAddr) = unaryOp epAddr L.epAddressToContract
compileExpr (MakeView e1 e2) = binaryOp e1 e2 (L.pair # L.wrapView)
compileExpr (MakeVoid e1 e2) = binaryOp e1 e2 (L.pair # L.wrapVoid)

compileExpr (Mem k c) = binaryOp k c L.mem
compileExpr (Size s) = unaryOp s L.size

compileExpr (StInsertNew l err k v store) =
  ternaryOp k v store $ L.stInsertNew l (failUsing err)
compileExpr (StInsert l k v store) =
  ternaryOp k v store $ L.stInsert l
compileExpr (StGet l ekey estore) = binaryOp ekey estore (L.stGet l)
compileExpr (StMem l ekey estore) = binaryOp ekey estore (L.stMem l)
compileExpr (StUpdate l ekey evalue estore) = ternaryOp ekey evalue estore (L.stUpdate l)
compileExpr (StDelete l ekey estore) = binaryOp ekey estore (L.stDelete l)

compileExpr (Wrap l exFld) = unaryOp exFld $ L.wrapOne l
compileExpr (Unwrap l exDt) = unaryOp exDt $ L.unwrapUnsafe_ l

compileExpr (ObjMan fldAcc) = compileObjectManipulation fldAcc
compileExpr (Construct _ fields) = IndigoState $ \md ->
  let cd = L.construct $ rmap (\e -> fieldCtor $ gcCode $ runIndigoState (compileExpr e) md) fields in
  GenCode (pushNoRef $ mdStack md) cd L.drop
compileExpr (ConstructWithoutNamed _ fields) = IndigoState $ \md ->
  let fieldCtrs =
          castFieldConstructors @a $
            rmap (fieldCtor . gcCode . usingIndigoState md . compileExpr) fields
  in GenCode (pushNoRef $ mdStack md) (L.construct @a fieldCtrs) L.drop
compileExpr (Name l e) = unaryOp e (toNamed l)
compileExpr (UnName l e) = unaryOp e (fromNamed l)

compileExpr (Slice ex1 ex2 ex3) = ternaryOp ex1 ex2 ex3 L.slice
compileExpr (Cast ex) = unaryOp ex L.cast
compileExpr (Concat ex1 ex2) = binaryOp ex1 ex2 L.concat
compileExpr (Concat' ex) = unaryOp ex L.concat'

compileExpr (ImplicitAccount kh) = unaryOp kh L.implicitAccount
compileExpr Now = nullaryOp L.now
compileExpr Sender = nullaryOp L.sender
compileExpr Amount = nullaryOp L.amount
compileExpr (CheckSignature pk sig bs) = ternaryOp pk sig bs L.checkSignature
compileExpr (Sha256 c) = unaryOp c L.sha256
compileExpr (Sha512 c) = unaryOp c L.sha512
compileExpr (Blake2b c) = unaryOp c L.blake2B
compileExpr (HashKey hk) = unaryOp hk L.hashKey
compileExpr ChainId = nullaryOp L.chainId
compileExpr Balance = nullaryOp L.balance

compileExpr EmptySet = nullaryOp L.emptySet

compileExpr (Get k m) = binaryOp k m L.get
compileExpr EmptyMap = nullaryOp L.emptyMap
compileExpr EmptyBigMap = nullaryOp L.emptyBigMap

compileExpr (Exec inp lambda) = binaryOp inp lambda L.exec
compileExpr (NonZero e) = unaryOp e L.nonZero

--------------------------------------------
-- Object manipulation: set, get fields
--------------------------------------------

-- | Compile 'ObjectManipulation' datatype to a cell on the stack.
-- This function leverages 'ObjManipulationRes' to put off actual field compilation.
compileObjectManipulation :: ObjectManipulation a -> IndigoState inp (a : inp)
compileObjectManipulation fa = IndigoState $ \md -> case runObjectManipulation (mdObjects md) fa of
  StillObject composite -> usingIndigoState md $ compileObjectF unNamedFieldExpr composite
  OnStack computation   -> usingIndigoState md computation

namedToExpr :: NamedFieldObj x name -> Expr (GetFieldType x name)
namedToExpr (NamedFieldObj flObj) = objToExpr namedToExpr flObj

-- | Convert arbitrary 'IndigoObjectF' into 'Expr'
-- with respect to given converter for fields.
objToExpr
  :: forall a f .
     (forall name . f name -> Expr (GetFieldType a name))
  -> IndigoObjectF f a
  -> Expr a
objToExpr _ (Cell refId) = V (Var @a refId)
objToExpr convExpr (Decomposed fields) =
  ConstructWithoutNamed (Proxy @a) (namedToTypedRec @a convExpr fields)

-- | Compile 'IndigoObjectF' to a stack cell,
-- with respect to given function that compiles inner fields.
compileObjectF
  :: forall a inp f .
     (forall name . f name -> Expr (GetFieldType a name))
  -> IndigoObjectF f a
  -> IndigoState inp (a : inp)
compileObjectF _ (Cell ref) = IndigoState $ \(mdStack -> s) ->
  GenCode (pushNoRef s) (varActionGet @a ref s) L.drop
compileObjectF conv obj = compileExpr $ objToExpr conv obj

-- | 'ObjManipulationRes' represents a postponed compilation of
-- 'ObjectManipulation' datatype. When 'ObjectManipulation' is being compiled
-- we are trying to put off the generation of code for work with an object
-- because we can just go to a deeper field without its "materialization"
-- onto stack.
data ObjManipulationRes inp a where
  StillObject :: ObjectExpr a -> ObjManipulationRes inp a
  OnStack :: IndigoState inp (a : inp) -> ObjManipulationRes inp a

-- | This function might look cumbersome
-- but basically it either goes deeper to an inner field or generates Lorentz code.
runObjectManipulation :: DecomposedObjects -> ObjectManipulation x -> ObjManipulationRes inp x
runObjectManipulation objs (Object e) = exprToManRes objs e

runObjectManipulation objs (ToField (v :: ObjectManipulation dt) (targetLb :: Label fname)) =
  case runObjectManipulation objs v of
    -- In case of decomposed fields, we just go deeper.
    StillObject (Decomposed fields) ->
      case fieldLens @dt @fname of
        -- If we access direct field, we just fetch it from fields
        TargetField lb _ -> exprToManRes objs $ unNamedFieldExpr (fetchField @dt lb fields)
        -- If we access deeper field, we fetch direct field and goes to the deeper field
        DeeperField lb _ ->
          let fe = unNamedFieldExpr $ fetchField @dt lb fields in
          runObjectManipulation objs (ToField (Object fe) targetLb)
    -- If stored object as cell on the stack, we get its field
    -- using 'sopToField', and since this moment 'ObjManipulationRes becomes
    -- a computation, not object anymore.
    StillObject (Cell refId) ->
      OnStack $ unaryOp (V $ Var refId) (sopToField @dt (flSFO fieldLens) targetLb)
    -- If we already got into computation, we use 'sopToField' to fetch field.
    OnStack compLHS -> OnStack $ IndigoState $ \mdI ->
      let cd = gcCode $ usingIndigoState mdI compLHS in
      GenCode (pushNoRef $ mdStack mdI) (cd # sopToField (flSFO fieldLens) targetLb) L.drop

runObjectManipulation objs (SetField (ev :: ObjectManipulation dt) (targetLb :: Label fname) ef) =
  case runObjectManipulation objs ev of
    StillObject lhsObj@(Decomposed fields) ->
      case fieldLens @dt @fname of
        -- If we set direct field, we just reassign its value with new one.
        TargetField lb _ ->
          StillObject $ Decomposed $ assignField @dt lb (NamedFieldExpr ef) fields
        -- If we set deeper field, we need to call recursively
        -- from a direct field, and set a target field of direct field.
        -- Getting a new value of direct field, we set the direct field to this value.
        DeeperField (lb :: Label interm) _ ->
          let fe = unNamedFieldExpr (fetchField @dt lb fields) in
          -- Computing new value of direct field
          case runObjectManipulation objs (SetField (Object fe) targetLb ef) of
            -- If it's still an object, we just reassign direct field with it.
            StillObject updField -> StillObject $ Decomposed $
              assignField @dt lb (NamedFieldExpr $ objToExpr unNamedFieldExpr updField) fields
            -- Otherwise, we use power of 'L.setField' to set a new value.
            OnStack rhs ->
              setFieldOnStack (compileObjectF unNamedFieldExpr lhsObj) rhs (L.setField @dt @interm lb)
    -- If stored object is Cell on stack, we set its field
    -- using 'sopSetField', and since this moment 'ObjManipulationRes' becomes
    -- a computation, not object anymore.
    StillObject (Cell refId) ->
      OnStack $ binaryOp ef (V $ Var refId) $ sopSetField (flSFO fieldLens) targetLb
    -- If we already got into computation, we use 'sopSetField' to set a field.
    OnStack compLHS ->
      setFieldOnStack compLHS (compileExpr ef) (sopSetField (flSFO $ fieldLens @dt) targetLb)
  where
    setFieldOnStack
      :: IndigoState inp (dt : inp)
      -> IndigoState (dt : inp) (fld : dt : inp)
      -> fld : dt : inp :-> dt : inp
      -> ObjManipulationRes inp dt
    setFieldOnStack lhs rhs setOp = OnStack $ IndigoState $ \mdI ->
      let GenCode st1 cdObj _cl1 = runIndigoState lhs mdI in
      let GenCode _st2 cdFld _cl2 = runIndigoState rhs (replStkMd mdI st1) in
      GenCode (pushNoRef $ mdStack mdI) (cdObj # cdFld # setOp) L.drop

-- | Convert an expression to 'ObjManipulationRes'.
-- The function pattern matches on some specific cases
-- of expression those compilation into a stack cell may be postponed.
-- They include 'Decomposed' variables and 'ConstructWithoutNamed' expressions.
--
-- This function can't be called for 'ObjMan' constructor, but we
-- take care of it just in case.
exprToManRes :: forall x inp . DecomposedObjects -> Expr x -> ObjManipulationRes inp x
exprToManRes objs (ObjMan objMan) = runObjectManipulation objs objMan
exprToManRes _ (ConstructWithoutNamed _ fields) =
  StillObject $ Decomposed $ typedToNamedRec @x NamedFieldExpr fields
exprToManRes objs (V var) = withObject objs var $ \case
  Cell refId ->
    StillObject $ Cell refId
  Decomposed fields ->
    StillObject $ Decomposed $ rmap (NamedFieldExpr . namedToExpr) fields
exprToManRes _ ex = OnStack $ compileExpr ex

---------------------------------------------------
-- Convenient helpers for operators compilation
---------------------------------------------------

ternaryOp
  :: KnownValue res
  => Expr n
  -> Expr m
  -> Expr l
  -> n : m : l : inp :-> res : inp
  -> IndigoState inp (res : inp)
ternaryOp e1 e2 e3 opCode = IndigoState $ \md ->
  let GenCode st3 cd3 _cl3  = runIndigoState (compileExpr e3) md in
  let GenCode st2 cd2 _cl2  = runIndigoState (compileExpr e2) (replStkMd md st3) in
  let GenCode _st1 cd1 _cl1 = runIndigoState (compileExpr e1) (replStkMd md st2) in
  GenCode (pushNoRef $ mdStack md) (cd3 # cd2 # cd1 # opCode) L.drop

binaryOp
  :: KnownValue res
  => Expr n -> Expr m
  -> n : m : inp :-> res : inp
  -> IndigoState inp (res : inp)
binaryOp e1 e2 opCode = IndigoState $ \md ->
  let GenCode st2 cd2 _cl2  = runIndigoState (compileExpr e2) md in
  let GenCode _st1 cd1 _cl1 = runIndigoState (compileExpr e1) (replStkMd md st2) in
  GenCode (pushNoRef $ mdStack md) (cd2 # cd1 # opCode) L.drop

unaryOp
  :: KnownValue res
  => Expr n
  -> n : inp :-> res : inp
  -> IndigoState inp (res : inp)
unaryOp e opCode = IndigoState $ \md ->
  let cd = gcCode $ runIndigoState (compileExpr e) md in
  GenCode (pushNoRef $ mdStack md) (cd # opCode) L.drop

nullaryOp :: KnownValue res => inp :-> res ': inp -> IndigoState inp (res ': inp)
nullaryOp lorentzInstr = IndigoState $ \md ->
  GenCode (pushNoRef $ mdStack md) lorentzInstr L.drop

ternaryOpFlat
  :: Expr n
  -> Expr m
  -> Expr l
  -> n : m : l : inp :-> inp
  -> IndigoState inp inp
ternaryOpFlat e1 e2 e3 opCode = IndigoState $ \md ->
  let GenCode st3 cd3 _cl3  = runIndigoState (compileExpr e3) md in
  let GenCode st2 cd2 _cl2  = runIndigoState (compileExpr e2) (replStkMd md st3) in
  let GenCode _st1 cd1 _cl1 = runIndigoState (compileExpr e1) (replStkMd md st2) in
  GenCode (mdStack md) (cd3 # cd2 # cd1 # opCode) L.nop

binaryOpFlat
  :: Expr n -> Expr m
  -> n : m : inp :-> inp
  -> IndigoState inp inp
binaryOpFlat e1 e2 opCode = IndigoState $ \md ->
  let GenCode st2 cd2 _cl2  = runIndigoState (compileExpr e2) md in
  let GenCode _st1 cd1 _cl1 = runIndigoState (compileExpr e1) (replStkMd md st2) in
  GenCode (mdStack md) (cd2 # cd1 # opCode) L.nop

unaryOpFlat
  :: Expr n
  -> n : inp :-> inp
  -> IndigoState inp inp
unaryOpFlat e opCode = IndigoState $ \md ->
  let cd = gcCode $ runIndigoState (compileExpr e) md in
  GenCode (mdStack md) (cd # opCode) L.nop

nullaryOpFlat :: inp :-> inp -> IndigoState inp inp
nullaryOpFlat lorentzInstr = IndigoState $ \md -> GenCode (mdStack md) lorentzInstr L.nop