Lookup.hs

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

{- |
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.Backend.Lookup
  ( -- * Variable Lookup Actions
    varActionGet
  , varActionSet
  , varActionUpdate
  , varActionOperation

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

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

import Indigo.Common.Var
  (HasSideEffects, Ops, RefId, StackVars(..), StackVars', StkEl(..), Var(..), operationsVar)
import Indigo.Lorentz
import Lorentz.Instr qualified as L
import Lorentz.Macro qualified as L
import Morley.Michelson.Typed (ToTs)
import Morley.Michelson.Typed.Instr.Constraints qualified as MI
import Morley.Michelson.Typed.T qualified as MI
import Morley.Util.Peano
import Morley.Util.Sing

----------------------------------------------------------------------------
-- 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 _ FailureStack = error "You try to get a cell on failure stack"
varActionGet ref (StkElements 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 _ FailureStack = error "You try to set a cell on failure stack"
varActionSet ref (StkElements 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 _ FailureStack _ = error "You try to update a cell on failure stack"
varActionUpdate v (StkElements 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
    Var 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).
     ( HasCallStack
     , KnownValue a
     )
  => Sing m -> RefId -> StackVars' inp -> inp :-> (a ': inp)
duupXVar m v stk =
  withVarMaybeDict (dupNLorentzConstraint n stk a out) $
  withDict (dupableConstraint a) $
    L.dupNPeano @('S m) @a @inp @(a ': inp)
  where
    a = Ref @a v
    n = SS m
    out = a :& 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 :: 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 $
    "You are looking for manually created or leaked variable. " <>
    pretty refId <> " of type " <> show (typeRep (Proxy @a))
  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. " <> pretty ref <>
      ".\nWas looking for a " <> show (typeRep $ Proxy @a) <>
      ", but found a: " <> show (typeRep $ Proxy @x)
  _ -> case varDepth @a ref xs of
    VarDepth idx -> VarDepth (SS idx)

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

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 $ withSingI m $ 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 $ withSingI m $ Dict

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

dupableConstraint
  :: (KnownValue a, HasCallStack)
  => StkEl a
  -> Dict (Dupable a)
dupableConstraint (_ :: StkEl a) =
  case decideOnDupable @a of
    IsDupable -> Dict
    IsNotDupable ->
      error "Non-dupable values (like tickets) are not supported yet"

dupNLorentzConstraint
  :: Sing n -> StackVars' inp -> StkEl a -> StackVars' out
  -> Maybe (Dict (L.ConstraintDUPNLorentz n inp out a))
dupNLorentzConstraint n inp a out = do
  Dict <- dupNConstraint n inp out a
  Dict <- dupNConstraint n (toTVals inp) (toTVals out) (toTVal a)
  return $ withSingI n $ Dict

updateNLorentzConstraint
  :: 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
  :: 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
  :: 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 $ withSingI n $ Dict

dipNLorentzConstraint
  :: 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 $ withSingI n $ Dict

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

dupNConstraint
  :: TestEquality any
  => Sing n -> Rec any inp -> Rec any out -> any a
  -> Maybe (Dict (MI.ConstraintDUPN' kind n inp out a))
dupNConstraint n inp out a = do
  Dict <- requireLongerOrSameLength inp n
  Dict <- isGreaterThan n SZ
  decN <- peanoSingDecrement n
  Refl <- testEquality inp ((lazyRtake decN inp) <+> a :& rdrop n inp)
  Refl <- testEquality out (a :& inp)
  return Dict

dugConstraint
  :: 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 (lazyRtake n out <+> rdrop (SS n) out) (rdrop (SS SZ) inp)
  Refl <- testEquality (lazyRtake n (rdrop (SS SZ) inp) <+> a :& rdrop n (rdrop (SS SZ) inp)) out
  return Dict

dipNConstraint
  :: 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 (lazyRtake n inp <+> s) inp
  Refl <- testEquality (lazyRtake n inp <+> s') out
  Refl <- testEquality (rdrop n inp) s
  Refl <- testEquality (rdrop n out) s'
  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 'SingI' constraint
data TVal (a :: MI.T) where
  TVal :: SingI a => TVal a

instance TestEquality TVal where
  testEquality (TVal :: TVal a) (TVal :: TVal b) = eqI @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"

lazyRtake :: Sing n -> Rec any s -> Rec any (LazyTake n s)
lazyRtake sn stk = case (sn, stk) of
  (SZ, _)           -> RNil
  (SS n, (x :& xs)) -> x :& lazyRtake 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"