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

{-# OPTIONS_GHC -Wno-orphans #-}

-- | Stack zipping.
--
-- This module provides functions for flattening stacks into tuples.
--
-- Also here we define an instance which turns any instruction,
-- not only lambdas, into a valid value.
module Lorentz.Zip
  ( ZipInstr (..)
  , zipInstr
  , unzipInstr
  , ZipInstrs
  , zippingStack
  , unzippingStack
  ) where

import Prelude hiding (drop)

import qualified Data.Kind as Kind

import Lorentz.Annotation
import Lorentz.Base
import Michelson.Typed
import Michelson.Untyped (noAnn)

-- | Version of '#' which performs some optimizations immediately.
--
-- In particular, this avoids glueing @Nop@s and @DIP Nop@s.
(##) :: (a :-> b) -> (b :-> c) -> (a :-> c)
I l :: Instr (ToTs a) (ToTs b)
l ## :: (a :-> b) -> (b :-> c) -> a :-> c
## I r :: Instr (ToTs b) (ToTs c)
r = Instr (ToTs a) (ToTs c) -> a :-> c
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (Instr (ToTs a) (ToTs b)
-> Instr (ToTs b) (ToTs c) -> Instr (ToTs a) (ToTs c)
forall (a :: [T]) (b :: [T]) (c :: [T]).
Instr a b -> Instr b c -> Instr a c
seqOpt Instr (ToTs a) (ToTs b)
l Instr (ToTs b) (ToTs c)
r)
l :: a :-> b
l ## r :: b :-> c
r = a :-> b
l (a :-> b) -> (b :-> c) -> a :-> c
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
# b :-> c
r

seqOpt :: Instr a b -> Instr b c -> Instr a c
seqOpt :: Instr a b -> Instr b c -> Instr a c
seqOpt l :: Instr a b
l r :: Instr b c
r =
  -- We are very verbose about cases to avoid
  -- significant compilation time increase
  case Instr a b
l of
    Nop -> Instr a c
Instr b c
r
    DIP Nop -> Instr a c
Instr b c
r
    x :: Instr a b
x -> case Instr b c
r of
      Nop -> Instr a b
Instr a c
x
      DIP Nop -> Instr a b
Instr a c
x
      _ -> Instr a b
l Instr a b -> Instr b c -> Instr a c
forall (a :: [T]) (b :: [T]) (c :: [T]).
Instr a b -> Instr b c -> Instr a c
`Seq` Instr b c
r

{- We have to work on low level because even "Lorentz.Instr" depends
   on this module.
-}

-- | Zipping stack into tuple and back.
class (KnownIsoT (ZippedStack s)) => ZipInstr (s :: [Kind.Type]) where
  -- | A type which contains the whole stack zipped.
  type ZippedStack s :: Kind.Type

  -- | Fold given stack into single value in typed Michelson.
  zipInstrTyped :: Instr (ToTs s) '[ToT (ZippedStack s)]

  -- | Unfold given stack from a single value in typed Michelson.
  unzipInstrTyped :: Instr '[ToT (ZippedStack s)] (ToTs s)


-- | Fold given stack into single value.
zipInstr :: forall s. ZipInstr s => s :-> '[ZippedStack s]
zipInstr :: s :-> '[ZippedStack s]
zipInstr = Instr (ToTs s) (ToTs '[ZippedStack s]) -> s :-> '[ZippedStack s]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (ZipInstr s => Instr (ToTs s) '[ToT (ZippedStack s)]
forall (s :: [*]).
ZipInstr s =>
Instr (ToTs s) '[ToT (ZippedStack s)]
zipInstrTyped @s)

-- | Unfold given stack from a single value.
unzipInstr :: forall s. ZipInstr s => '[ZippedStack s] :-> s
unzipInstr :: '[ZippedStack s] :-> s
unzipInstr = Instr (ToTs '[ZippedStack s]) (ToTs s) -> '[ZippedStack s] :-> s
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (ZipInstr s => Instr '[ToT (ZippedStack s)] (ToTs s)
forall (s :: [*]).
ZipInstr s =>
Instr '[ToT (ZippedStack s)] (ToTs s)
unzipInstrTyped @s)

instance ZipInstr '[] where
  type ZippedStack '[] = ()
  zipInstrTyped :: Instr (ToTs '[]) '[ToT (ZippedStack '[])]
zipInstrTyped = Instr (ToTs '[]) '[ToT (ZippedStack '[])]
forall (inp :: [T]). Instr inp ('TUnit : inp)
UNIT
  unzipInstrTyped :: Instr '[ToT (ZippedStack '[])] (ToTs '[])
unzipInstrTyped = Instr '[ToT (ZippedStack '[])] (ToTs '[])
forall (a :: T) (out :: [T]). Instr (a : out) out
DROP

instance (KnownIsoT a) => ZipInstr '[a] where
  type ZippedStack '[a] = a
  zipInstrTyped :: Instr (ToTs '[a]) '[ToT (ZippedStack '[a])]
zipInstrTyped = Instr (ToTs '[a]) '[ToT (ZippedStack '[a])]
forall (inp :: [T]). Instr inp inp
Nop
  unzipInstrTyped :: Instr '[ToT (ZippedStack '[a])] (ToTs '[a])
unzipInstrTyped = Instr '[ToT (ZippedStack '[a])] (ToTs '[a])
forall (inp :: [T]). Instr inp inp
Nop

-- | Such definition seems the only possible one we can support
-- efficiently.
instance (ZipInstr (b ': s), KnownIsoT a) => ZipInstr (a ': b ': s) where
  type ZippedStack (a ': b ': s) = (a, ZippedStack (b ': s))
  zipInstrTyped :: Instr (ToTs (a : b : s)) '[ToT (ZippedStack (a : b : s))]
zipInstrTyped = Instr (ToT b : ToTs s) '[ToT (ZippedStack (b : s))]
-> Instr
     (ToT a : ToT b : ToTs s) '[ToT a, ToT (ZippedStack (b : s))]
forall (a :: [T]) (c :: [T]) (b :: T).
Instr a c -> Instr (b : a) (b : c)
DIP (ZipInstr (b : s) =>
Instr (ToTs (b : s)) '[ToT (ZippedStack (b : s))]
forall (s :: [*]).
ZipInstr s =>
Instr (ToTs s) '[ToT (ZippedStack s)]
zipInstrTyped @(b ': s)) Instr (ToT a : ToT b : ToTs s) '[ToT a, ToT (ZippedStack (b : s))]
-> Instr
     '[ToT a, ToT (ZippedStack (b : s))]
     '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
-> Instr
     (ToT a : ToT b : ToTs s)
     '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
forall (a :: [T]) (b :: [T]) (c :: [T]).
Instr a b -> Instr b c -> Instr a c
`seqOpt` Instr
  '[ToT a, ToT (ZippedStack (b : s))]
  '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
forall (i :: [T]) (o :: [T]) (a :: T) (b :: T) (s :: [T]).
(i ~ (a : b : s), o ~ ('TPair a b : s)) =>
Instr i o
PAIR
  unzipInstrTyped :: Instr '[ToT (ZippedStack (a : b : s))] (ToTs (a : b : s))
unzipInstrTyped = Instr
  '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
  '[ToT a, ToT (ZippedStack (b : s))]
forall (i :: [T]) (o :: [T]) (a :: T) (b :: T) (s :: [T]).
(i ~ ('TPair a b : s), o ~ (a : b : s)) =>
Instr i o
UNPAIR Instr
  '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
  '[ToT a, ToT (ZippedStack (b : s))]
-> Instr
     '[ToT a, ToT (ZippedStack (b : s))] (ToT a : ToT b : ToTs s)
-> Instr
     '[ 'TPair (ToT a) (ToT (ZippedStack (b : s)))]
     (ToT a : ToT b : ToTs s)
forall (a :: [T]) (b :: [T]) (c :: [T]).
Instr a b -> Instr b c -> Instr a c
`seqOpt` Instr '[ToT (ZippedStack (b : s))] (ToT b : ToTs s)
-> Instr
     '[ToT a, ToT (ZippedStack (b : s))] (ToT a : ToT b : ToTs s)
forall (a :: [T]) (c :: [T]) (b :: T).
Instr a c -> Instr (b : a) (b : c)
DIP (ZipInstr (b : s) =>
Instr '[ToT (ZippedStack (b : s))] (ToTs (b : s))
forall (s :: [*]).
ZipInstr s =>
Instr '[ToT (ZippedStack s)] (ToTs s)
unzipInstrTyped @(b ': s))

-- | Require several stacks to comply 'ZipInstr' constraint.
type ZipInstrs ss = Each '[ZipInstr] ss

-- | Flatten both ends of instruction stack.
zippingStack
  :: ZipInstrs [inp, out]
  => inp :-> out -> Lambda (ZippedStack inp) (ZippedStack out)
zippingStack :: (inp :-> out) -> Lambda (ZippedStack inp) (ZippedStack out)
zippingStack code :: inp :-> out
code = '[ZippedStack inp] :-> inp
forall (s :: [*]). ZipInstr s => '[ZippedStack s] :-> s
unzipInstr ('[ZippedStack inp] :-> inp)
-> (inp :-> out) -> '[ZippedStack inp] :-> out
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## inp :-> out
code ('[ZippedStack inp] :-> out)
-> (out :-> '[ZippedStack out])
-> Lambda (ZippedStack inp) (ZippedStack out)
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## out :-> '[ZippedStack out]
forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr

-- | Unflatten both ends of instruction stack.
unzippingStack
  :: ZipInstrs [inp, out]
  => Lambda (ZippedStack inp) (ZippedStack out) -> inp :-> out
unzippingStack :: Lambda (ZippedStack inp) (ZippedStack out) -> inp :-> out
unzippingStack code :: Lambda (ZippedStack inp) (ZippedStack out)
code = inp :-> '[ZippedStack inp]
forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr (inp :-> '[ZippedStack inp])
-> Lambda (ZippedStack inp) (ZippedStack out)
-> inp :-> '[ZippedStack out]
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## Lambda (ZippedStack inp) (ZippedStack out)
code (inp :-> '[ZippedStack out])
-> ('[ZippedStack out] :-> out) -> inp :-> out
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## '[ZippedStack out] :-> out
forall (s :: [*]). ZipInstr s => '[ZippedStack s] :-> s
unzipInstr

instance (WellTypedToT (ZippedStack inp), WellTypedToT (ZippedStack out), ZipInstr inp, ZipInstr out) => IsoValue (inp :-> out) where
  type ToT (inp :-> out) = 'TLambda (ToT (ZippedStack inp)) (ToT (ZippedStack out))
  toVal :: (inp :-> out) -> Value (ToT (inp :-> out))
toVal i :: inp :-> out
i = RemFail Instr '[ToT (ZippedStack inp)] '[ToT (ZippedStack out)]
-> Value'
     Instr ('TLambda (ToT (ZippedStack inp)) (ToT (ZippedStack out)))
forall (inp :: T) (out :: T) (instr :: [T] -> [T] -> *).
(KnownT inp, KnownT out,
 forall (i :: [T]) (o :: [T]). Show (instr i o),
 forall (i :: [T]) (o :: [T]). Eq (instr i o),
 forall (i :: [T]) (o :: [T]). NFData (instr i o)) =>
RemFail instr '[inp] '[out] -> Value' instr ('TLambda inp out)
VLam (RemFail Instr '[ToT (ZippedStack inp)] '[ToT (ZippedStack out)]
 -> Value'
      Instr ('TLambda (ToT (ZippedStack inp)) (ToT (ZippedStack out))))
-> (('[ZippedStack inp] :-> '[ZippedStack out])
    -> RemFail Instr '[ToT (ZippedStack inp)] '[ToT (ZippedStack out)])
-> ('[ZippedStack inp] :-> '[ZippedStack out])
-> Value'
     Instr ('TLambda (ToT (ZippedStack inp)) (ToT (ZippedStack out)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ('[ZippedStack inp] :-> '[ZippedStack out])
-> RemFail Instr '[ToT (ZippedStack inp)] '[ToT (ZippedStack out)]
forall (inp :: [*]) (out :: [*]).
(inp :-> out) -> RemFail Instr (ToTs inp) (ToTs out)
unLorentzInstr (('[ZippedStack inp] :-> '[ZippedStack out])
 -> Value (ToT (inp :-> out)))
-> ('[ZippedStack inp] :-> '[ZippedStack out])
-> Value (ToT (inp :-> out))
forall a b. (a -> b) -> a -> b
$ (inp :-> out) -> '[ZippedStack inp] :-> '[ZippedStack out]
forall (inp :: [*]) (out :: [*]).
ZipInstrs '[inp, out] =>
(inp :-> out) -> Lambda (ZippedStack inp) (ZippedStack out)
zippingStack inp :-> out
i
  fromVal :: Value (ToT (inp :-> out)) -> inp :-> out
fromVal (VLam i :: RemFail Instr '[inp] '[out]
i) = inp :-> '[ZippedStack inp]
forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr (inp :-> '[ZippedStack inp])
-> ('[ZippedStack inp] :-> '[ZippedStack out])
-> inp :-> '[ZippedStack out]
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## RemFail Instr (ToTs '[ZippedStack inp]) (ToTs '[ZippedStack out])
-> '[ZippedStack inp] :-> '[ZippedStack out]
forall (inp :: [*]) (out :: [*]).
RemFail Instr (ToTs inp) (ToTs out) -> inp :-> out
LorentzInstr RemFail Instr '[inp] '[out]
RemFail Instr (ToTs '[ZippedStack inp]) (ToTs '[ZippedStack out])
i (inp :-> '[ZippedStack out])
-> ('[ZippedStack out] :-> out) -> inp :-> out
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## '[ZippedStack out] :-> out
forall (s :: [*]). ZipInstr s => '[ZippedStack s] :-> s
unzipInstr

instance
    ( HasAnnotation (ZippedStack i)
    , HasAnnotation (ZippedStack o)
    )
  =>
    HasAnnotation (i :-> o)
  where
  getAnnotation :: FollowEntrypointFlag -> Notes (ToT (i :-> o))
getAnnotation b :: FollowEntrypointFlag
b = TypeAnn
-> Notes (ToT (ZippedStack i))
-> Notes (ToT (ZippedStack o))
-> Notes ('TLambda (ToT (ZippedStack i)) (ToT (ZippedStack o)))
forall (p :: T) (q :: T).
TypeAnn -> Notes p -> Notes q -> Notes ('TLambda p q)
NTLambda TypeAnn
forall k (a :: k). Annotation a
noAnn
    (FollowEntrypointFlag -> Notes (ToT (ZippedStack i))
forall a. HasAnnotation a => FollowEntrypointFlag -> Notes (ToT a)
getAnnotation @(ZippedStack i) FollowEntrypointFlag
b)
    (FollowEntrypointFlag -> Notes (ToT (ZippedStack o))
forall a. HasAnnotation a => FollowEntrypointFlag -> Notes (ToT a)
getAnnotation @(ZippedStack o) FollowEntrypointFlag
b)