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

-- | 'Expr' compilation

module Indigo.Internal.Expr.Compilation
  ( compileExpr

  , ObjManipulationRes (..)
  , runObjectManipulation

  , 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 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(..), NamedFieldVar(..), castFieldConstructors, namedToTypedRec, pushNoRefMd,
  typedToNamedRec)
import Indigo.Internal.State
  (GenCode(..), IndigoState(..), MetaData(..), iget, iput, usingIndigoState, (>>=))
import Indigo.Lorentz

compileExpr :: forall a inp . Expr a -> IndigoState inp (a & inp) ()
compileExpr (C a) = do
  md <- iget
  iput $ GenCode () (pushNoRefMd md) (L.push a) L.drop
compileExpr (V v) = compileObjectF (\(NamedFieldVar fl) -> V fl) v
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 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 (UInsertNew l err k v store) =
  ternaryOp k v store $ ustoreInsertNew l (failUsing err)
compileExpr (UInsert l k v store) =
  ternaryOp k v store $ ustoreInsert l
compileExpr (UGet l ekey estore) = binaryOp ekey estore (ustoreGet l)
compileExpr (UMem l ekey estore) = binaryOp ekey estore (ustoreMem l)
compileExpr (UUpdate l ekey evalue estore) = ternaryOp ekey evalue estore (ustoreUpdate l)
compileExpr (UDelete l ekey estore) = binaryOp ekey estore (ustoreDelete 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 () (pushNoRefMd md) cd L.drop
compileExpr (ConstructWithoutNamed fields) = IndigoState $ \md ->
  let fieldCtrs =
          castFieldConstructors @a $
            rmap (fieldCtor . gcCode . usingIndigoState md . compileExpr) fields
  in GenCode () (pushNoRefMd 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

-- | Convert arbitrary 'IndigoObjectF' into 'Expr',
-- having converter for fields.
objToExpr
  :: forall a f .
     (forall name . f name -> Expr (GetFieldType a name))
  -> IndigoObjectF f a
  -> Expr a
objToExpr _ (Cell refId) = V (Cell @a refId)
objToExpr convExpr (Decomposed fields) =
  ConstructWithoutNamed $ namedToTypedRec @a convExpr fields

-- | Compile 'IndigoObjectF' to a stack cell,
-- having a function which 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) = do
  md@(MetaData s _) <- iget
  iput $ GenCode () (pushNoRefMd md) (varActionGet @a ref s) L.drop
compileObjectF conv obj = compileExpr $ objToExpr conv obj

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

-- | '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 it basically either goes deeper to an inner field or generates Lorentz code.
runObjectManipulation :: ObjectManipulation x -> ObjManipulationRes inp x
runObjectManipulation (Object e) = exprToManRes e

runObjectManipulation (ToField (v :: ObjectManipulation dt) (targetLb :: Label fname)) =
  case runObjectManipulation 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 $ 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 (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 $ Cell refId) (sopToField @dt (flSFO fieldLens) targetLb)
    -- If we already got into computation, we use 'sopToField' to fetch field.
    OnStack compLHS -> OnStack $ IndigoState $ \md ->
      let cd = gcCode $ runIndigoState compLHS md in
      GenCode () (pushNoRefMd md) (cd # sopToField (flSFO fieldLens) targetLb) L.drop

runObjectManipulation (SetField (ev :: ObjectManipulation dt) (targetLb :: Label fname) ef) =
  case runObjectManipulation 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 (SetField (Object fe) targetLb ef) of
            -- If it's still 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 $ Cell 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 $ \md ->
      let GenCode _ md1 cdObj _cl1 = runIndigoState lhs md in
      let GenCode _ _md2 cdFld _cl2 = runIndigoState rhs md1 in
      GenCode () (pushNoRefMd md) (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 . Expr x -> ObjManipulationRes inp x
exprToManRes (ObjMan objMan) = runObjectManipulation objMan
exprToManRes (ConstructWithoutNamed fields) =
  StillObject $ Decomposed $ typedToNamedRec @x NamedFieldExpr fields
exprToManRes (V (Decomposed fields)) =
  StillObject $ Decomposed $ rmap (\(NamedFieldVar f) -> NamedFieldExpr $ V f) fields
exprToManRes (V (Cell refId)) = StillObject $ Cell refId
exprToManRes ex = OnStack $ compileExpr ex

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 _ md3 cd3 _cl3  = runIndigoState (compileExpr e3) md in
  let GenCode _ md2 cd2 _cl2  = runIndigoState (compileExpr e2) md3 in
  let GenCode _ _md1 cd1 _cl1 = runIndigoState (compileExpr e1) md2 in
  GenCode () (pushNoRefMd 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 _ md2 cd2 _cl2  = runIndigoState (compileExpr e2) md in
  let GenCode _ _md1 cd1 _cl1 = runIndigoState (compileExpr e1) md2 in
  GenCode () (pushNoRefMd 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 () (pushNoRefMd md) (cd # opCode) L.drop

nullaryOp :: KnownValue res => inp :-> res ': inp -> IndigoState inp (res ': inp) ()
nullaryOp lorentzInstr = IndigoState $ \md ->
  GenCode () (pushNoRefMd 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 _ md3 cd3 _cl3  = runIndigoState (compileExpr e3) md in
  let GenCode _ md2 cd2 _cl2  = runIndigoState (compileExpr e2) md3 in
  let GenCode _ _md1 cd1 _cl1 = runIndigoState (compileExpr e1) md2 in
  GenCode () 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 _ md2 cd2 _cl2  = runIndigoState (compileExpr e2) md in
  let GenCode _ _md1 cd1 _cl1 = runIndigoState (compileExpr e1) md2 in
  GenCode () 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 () md (cd # opCode) L.nop

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