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

{- |
This module contains the logic to lookup 'Var's in a stack and the actions to
manipulate it.

For efficiency, actions are implemented using Lorentz macros.
To do so every necessary constraint is checked at runtime.
-}

module Indigo.Internal.Lookup
  ( -- * Variable Lookup Actions
    varActionGet
  , varActionSet
  , varActionUpdate
  , varActionOperation

  -- * Vinyl manipulation helpers
  , rtake
  , rdrop
  ) where

import Data.Constraint (Dict(..), HasDict)
import Data.Singletons (Sing, SingI(..))
import Data.Type.Equality (TestEquality(..))
import Data.Typeable ((:~:)(..), eqT)
import Data.Vinyl ((<+>))
import Data.Vinyl.TypeLevel (type (++))
import Prelude hiding (tail)

import Indigo.Internal.Object (HasSideEffects, IndigoObjectF(..), Ops, operationsVar)
import Indigo.Internal.State (RefId, StackVars, StkEl(..))
import Indigo.Lorentz
import qualified Lorentz.Instr as L
import qualified Lorentz.Macro as L
import Michelson.Typed (ToTs)
import qualified Michelson.Typed.Instr as MI
import qualified Michelson.Typed.T as MI
import Util.Peano

----------------------------------------------------------------------------
-- Variable Lookup Actions
----------------------------------------------------------------------------

-- | Puts a copy of the value for the given 'Var' on top of the stack
varActionGet :: forall a stk . KnownValue a => RefId -> StackVars stk -> stk :-> a & stk
varActionGet ref stk = case varDepth @a ref stk of
  VarDepth n -> duupXVar n ref stk

-- | Sets the value for the given 'Var' to the topmost value on the stack
varActionSet :: forall a stk . KnownValue a => RefId -> StackVars stk -> a & stk :-> stk
varActionSet ref stk = case varDepth @a ref stk of
  VarDepth n -> replaceNVar (SS n) ref (NoRef :& stk)

-- | Updates the value for the given 'Var' with the topmost value on the stack
-- using the given binary instruction.
varActionUpdate
  :: forall a b stk . (KnownValue a, KnownValue b)
  => RefId
  -> StackVars stk
  -> '[b, a] :-> '[a]
  -> (b ': stk) :-> stk
varActionUpdate v stk instr = case varDepth @a v stk of
  VarDepth n -> updateNVar (SS n) v (NoRef :& stk) instr

-- | Given a stack with a list of 'Operation's on its bottom, updates it by
-- appending the 'Operation' on the top.
varActionOperation
  :: HasSideEffects
  => StackVars stk
  -> (Operation ': stk) :-> stk
varActionOperation s =
  case operationsVar of
    Cell refId -> varActionUpdate @Ops refId s L.cons

----------------------------------------------------------------------------
-- Variable-based Macros
----------------------------------------------------------------------------

-- | Like 'varActionGet', but requires the depth of the 'Var' in the input stack
duupXVar
  :: forall inp a (m :: Peano) s tail.
     ( HasCallStack
     , KnownValue a, KnownPeano m, SingI m
     , s ~ ((Take m inp) ++ (a ':  a ': Drop ('S m) inp))
     , tail ~ (Drop ('S m) inp)
     )
  => Sing m -> RefId -> StackVars inp -> inp :-> (a ': inp)
duupXVar m v stk =
  withVarMaybeDict (duupXClassConstraint n stk a s tail) $
    L.duupXImpl @('S m) @inp @a @s @tail
  where
    a = Ref @a v
    n = SS m
    s = (rtake m stk) <+> (a :& a :& rdrop n stk)
    tail = rdrop n stk

-- | Like 'varActionSet', but requires the depth of the 'Var' in the input stack
replaceNVar
  :: forall s a (n :: Peano) mid tail.
     ( HasCallStack
     , KnownValue a
     , mid ~ (Take n (a ': s) ++ Drop n s)
     , tail ~ Drop n s
     )
  => Sing n -> RefId -> StackVars (a ': s) -> (a ': s) :-> s
replaceNVar n v stk@(_ :& s) =
  withVarMaybeDict (replaceNClassConstraint n s (Ref @a v) mid tail) $
    L.replaceNImpl @n @s @a @mid @tail
  where
    mid = rtake n stk <+> rdrop n s
    tail = rdrop n s

-- | Like 'varActionUpdate', but requires the depth of the 'Var' in the input stack
updateNVar
  :: forall s a b mid tail (n :: Peano).
     ( HasCallStack
     , KnownValue b
     , tail ~ Drop n s
     , mid ~ ((Drop ('S 'Z) (Take n (a ': s))) ++ (a ': Drop n (a ': s)))
     )
  => Sing n -> RefId -> StackVars (a ': s)
  -> '[a, b] :-> '[b]
  -> (a ': s) :-> s
updateNVar n v stk@(a :& s) instr =
  withVarMaybeDict (updateNClassConstraint n s a (Ref @b v) mid tail) $
    L.updateNImpl @n @s @a @b @mid @tail instr
  where
    mid = rdrop (SS SZ) (rtake n stk) <+> a :& rdrop n stk
    tail = rdrop n s

withVarMaybeDict :: (HasDict c e, HasCallStack) => Maybe e -> (c => r) -> r
withVarMaybeDict mDict = withDict (fromMaybe (error constraintFailure) mDict)
  where
    -- NOTE: provided that the 'VarDepth' is correctly calculated every place
    -- where this is used should never result in a 'Nothing', as this is only
    -- necessary to prove to GHC properties that always hold.
    -- For this reason a failure here is always unexpected
    constraintFailure = "Unexpected failure in Var's constraint checking"

----------------------------------------------------------------------------
-- Variable's Depth
----------------------------------------------------------------------------

-- | Keeps the 0-indexed depth of a variable as a 'Peano' 'Sing'leton
data VarDepth where
  VarDepth
    :: forall idx . (KnownPeano idx, SingI idx)
    => Sing (idx :: Peano) -> VarDepth

-- | Calculates the 'VarDepth' of the given 'Var' in the stack
varDepth
  :: forall a s. KnownValue a
  => RefId
  -> StackVars s
  -> VarDepth
varDepth refId = \case
  RNil -> error $ "Manually created or leaked variable. Ref #" <> show refId
  stk@(_ :& _) -> varDepthNonEmpty @a refId stk

varDepthNonEmpty
  :: forall a s x xs. (KnownValue a, s ~ (x & xs))
  => RefId -> StackVars s -> VarDepth
varDepthNonEmpty ref (x :& xs) = case x of
  Ref topRef | ref == topRef -> case eqT @a @x of
    Just Refl -> VarDepth SZ
    Nothing   -> error $ "Invalid variable type. Ref #" <> show ref
  _ -> case varDepth @a ref xs of
    VarDepth idx -> VarDepth (SS idx)

----------------------------------------------------------------------------
-- Macro class constraints
----------------------------------------------------------------------------

duupXClassConstraint
  :: Sing n -> StackVars inp -> StkEl a -> StackVars s -> StackVars s'
  -> Maybe (Dict (L.DuupX n inp a s s'))
duupXClassConstraint n inp a s s' = case n of
  SZ -> Nothing
  SS SZ -> case inp of
    (x :& _xs) -> do
      Refl <- testEquality x a
      return Dict
    _ -> Nothing
  SS (SS SZ) -> case inp of
    (_ :& x :& _xs) -> do
      Refl <- testEquality x a
      return Dict
    _ -> Nothing
  SS (SS (SS m)) -> do
    Dict <- duupXLorentzConstraint (SS (SS m)) inp a s s'
    return Dict

replaceNClassConstraint
  :: Sing n -> StackVars s -> StkEl a -> StackVars mid -> StackVars tail
  -> Maybe (Dict (L.ReplaceN n s a mid tail))
replaceNClassConstraint n s a mid tail = case n of
  SZ -> Nothing
  SS SZ -> case s of
    (x :& _xs) -> do
      Refl <- testEquality x a
      return Dict
    _ -> Nothing
  SS (SS m) -> do
    Dict <- replaceNLorentzConstraint (SS m) s a mid tail
    return Dict

updateNClassConstraint
  :: Sing n -> StackVars s -> StkEl a -> StkEl b -> StackVars mid -> StackVars tail
  -> Maybe (Dict (L.UpdateN n s a b mid tail))
updateNClassConstraint n s a b mid tail = case n of
  SZ -> Nothing
  SS SZ -> case s of
    (x :& xs) -> do
      Refl <- testEquality x b
      Refl <- testEquality xs tail
      return Dict
    _ -> Nothing
  SS (SS SZ) -> case s of
    (_x :& y :& xs) -> do
      Refl <- testEquality y b
      Refl <- testEquality xs tail
      return Dict
    _ -> Nothing
  SS (SS (SS m)) -> do
    Dict <- updateNLorentzConstraint (SS (SS m)) s a b mid tail
    return Dict

----------------------------------------------------------------------------
-- Lorentz constraints
----------------------------------------------------------------------------

duupXLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars inp -> StkEl a -> StackVars s -> StackVars s'
  -> Maybe (Dict (L.ConstraintDuupXLorentz n inp a s s'))
duupXLorentzConstraint n s a s1 tail = do
  Refl <- testEquality tail (rdrop (SS n) s)
  Refl <- testEquality (toTVals tail) (rdrop (SS n) (toTVals s))
  Dict <- dipNLorentzConstraint n s s1 (a :& tail) (a :& a :& tail)
  Dict <- digLorentzConstraint n s1 (a :& s) a
  return Dict

updateNLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars s -> StkEl a -> StkEl b -> StackVars mid -> StackVars tail
  -> Maybe (Dict (L.ConstraintUpdateNLorentz n s a b mid tail))
updateNLorentzConstraint n s a b mid tail = do
  Refl <- testEquality tail (rdrop (SS n) s)
  Refl <- testEquality (toTVals tail) (rdrop (SS n) (toTVals s))
  Dict <- dugLorentzConstraint n (a :& s) mid a
  Dict <- dipNLorentzConstraint n mid s (a :& b :& tail) (b :& tail)
  return Dict

replaceNLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars s -> StkEl a -> StackVars mid -> StackVars tail
  -> Maybe (Dict (L.ConstraintReplaceNLorentz n s a mid tail))
replaceNLorentzConstraint n s a mid tail = do
  Refl <- testEquality tail (rdrop (SS n) s)
  Refl <- testEquality (toTVals tail) (rdrop (SS n) (toTVals s))
  Dict <- dipNLorentzConstraint (SS n) (a :& s) mid (a :& tail) tail
  Dict <- dugLorentzConstraint n mid s a
  return Dict

dugLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars inp -> StackVars out -> StkEl a
  -> Maybe (Dict (L.ConstraintDUGLorentz n inp out a))
dugLorentzConstraint n inp out a = do
  Dict <- dugConstraint n inp out a
  Dict <- dugConstraint n (toTVals inp) (toTVals out) (toTVal a)
  return Dict

dipNLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars inp -> StackVars out -> StackVars s -> StackVars s'
  -> Maybe (Dict (L.ConstraintDIPNLorentz n inp out s s'))
dipNLorentzConstraint n inp out s s' = do
  Dict <- dipNConstraint n inp out s s'
  Dict <- dipNConstraint n (toTVals inp) (toTVals out) (toTVals s) (toTVals s')
  return Dict

digLorentzConstraint
  :: (KnownPeano n, SingI n)
  => Sing n -> StackVars inp -> StackVars out -> StkEl a
  -> Maybe (Dict (L.ConstraintDIGLorentz n inp out a))
digLorentzConstraint n inp out a = do
  Dict <- digConstraint n inp out a
  Dict <- digConstraint n (toTVals inp) (toTVals out) (toTVal a)
  return Dict

----------------------------------------------------------------------------
-- Morley constraints
----------------------------------------------------------------------------

dugConstraint
  :: (KnownPeano n, SingI n, TestEquality any)
  => Sing n -> Rec any inp -> Rec any out -> any a
  -> Maybe (Dict (MI.ConstraintDUG' kind n inp out a))
dugConstraint n inp out a = do
  Dict <- requireLongerThan out n
  Refl <- testEquality (a :& rdrop (SS SZ) inp) inp
  Refl <- testEquality (rtake n (rdrop (SS SZ) inp) <+> a :& rdrop (SS n) inp) out
  return Dict

dipNConstraint
  :: (KnownPeano n, SingI n, TestEquality any)
  => Sing n -> Rec any inp -> Rec any out -> Rec any s -> Rec any s'
  -> Maybe (Dict (MI.ConstraintDIPN' kind n inp out s s'))
dipNConstraint n inp out s s' = do
  Dict <- requireLongerOrSameLength inp n
  Refl <- testEquality (rtake n inp <+> s) inp
  Refl <- testEquality (rtake n inp <+> s') out
  return Dict

digConstraint
  :: (KnownPeano n, SingI n, TestEquality any)
  => Sing n -> Rec any inp -> Rec any out -> any a
  -> Maybe (Dict (MI.ConstraintDIG' kind n inp out a))
digConstraint n inp out a = do
  Dict <- requireLongerThan inp n
  Refl <- testEquality inp (rtake n inp <+> a :& rdrop (SS n) inp)
  Refl <- testEquality out (a :& rtake n inp <+> rdrop (SS n) inp)
  return Dict

----------------------------------------------------------------------------
-- Conversion for ToT constraints
----------------------------------------------------------------------------

-- | Stack representation of 'MI.T' used for constraint checking
type TValStack (stk :: [MI.T]) = Rec TVal stk

-- | Simple datatype used to keep a 'MI.T' and its 'Typeable' constraint
data TVal (a :: MI.T) where
  TVal :: Typeable a => TVal a

instance TestEquality TVal where
  testEquality (TVal :: TVal a) (TVal :: TVal b) = eqT @a @b

toTVal :: forall a. StkEl a -> TVal (ToT a)
toTVal = \case
  NoRef -> TVal @(ToT a)
  Ref _ -> TVal @(ToT a)

toTVals :: StackVars stk -> TValStack (ToTs stk)
toTVals = \case
  RNil          -> RNil
  (NoRef :& xs) -> TVal :& toTVals xs
  (Ref _ :& xs) -> TVal :& toTVals xs

----------------------------------------------------------------------------
-- Vinyl manipulation helpers
----------------------------------------------------------------------------

rtake :: Sing n -> Rec any s -> Rec any (Take n s)
rtake sn stk = case (sn, stk) of
  (SZ, _)           -> RNil
  (SS n, (x :& xs)) -> x :& rtake n xs
  (SS _, RNil)      -> error "given stack is too small"

rdrop :: Sing n -> Rec any s -> Rec any (Drop n s)
rdrop sn stk = case (sn, stk) of
  (SZ, _)           -> stk
  (SS n, (_ :& xs)) -> rdrop n xs
  (SS _, RNil)      -> error "given stack is too small"