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

-- | 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
  , ZippedStackRepr(..)
  , ZSNil(..)
  , (##)
  ) where

import Prelude hiding (drop)

import Fmt (Buildable(..), (+|), (|+))

import Lorentz.Annotation
import Lorentz.Base
import Morley.Michelson.Typed

-- | 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 Instr (ToTs a) (ToTs b)
l ## :: forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## I 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)
a :-> b
l ## 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 :: forall (a :: [T]) (b :: [T]) (c :: [T]).
Instr a b -> Instr b c -> Instr a c
seqOpt Instr a b
l Instr b c
r =
  -- We are very verbose about cases to avoid
  -- significant compilation time increase
  case Instr a b
l of
    Instr a b
Nop -> Instr a c
Instr b c
r
    DIP Instr a c
Nop -> Instr a c
Instr b c
r
    Instr a b
x -> case Instr b c
r of
      Instr b c
Nop -> Instr a b
Instr a c
x
      DIP Instr a c
Nop -> Instr a b
Instr a c
x
      Instr b c
_ -> 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.
-}

infixr 5 :::
-- | A type used to represent a zipped stack of at least two elements,
-- isomorphic to a pair and represented as such in Michelson.
data ZippedStackRepr a b = a ::: b
  deriving stock (Int -> ZippedStackRepr a b -> ShowS
[ZippedStackRepr a b] -> ShowS
ZippedStackRepr a b -> String
(Int -> ZippedStackRepr a b -> ShowS)
-> (ZippedStackRepr a b -> String)
-> ([ZippedStackRepr a b] -> ShowS)
-> Show (ZippedStackRepr a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> ZippedStackRepr a b -> ShowS
forall a b. (Show a, Show b) => [ZippedStackRepr a b] -> ShowS
forall a b. (Show a, Show b) => ZippedStackRepr a b -> String
showList :: [ZippedStackRepr a b] -> ShowS
$cshowList :: forall a b. (Show a, Show b) => [ZippedStackRepr a b] -> ShowS
show :: ZippedStackRepr a b -> String
$cshow :: forall a b. (Show a, Show b) => ZippedStackRepr a b -> String
showsPrec :: Int -> ZippedStackRepr a b -> ShowS
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> ZippedStackRepr a b -> ShowS
Show, ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
(ZippedStackRepr a b -> ZippedStackRepr a b -> Bool)
-> (ZippedStackRepr a b -> ZippedStackRepr a b -> Bool)
-> Eq (ZippedStackRepr a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b.
(Eq a, Eq b) =>
ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
/= :: ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
$c/= :: forall a b.
(Eq a, Eq b) =>
ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
== :: ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
$c== :: forall a b.
(Eq a, Eq b) =>
ZippedStackRepr a b -> ZippedStackRepr a b -> Bool
Eq, (forall x. ZippedStackRepr a b -> Rep (ZippedStackRepr a b) x)
-> (forall x. Rep (ZippedStackRepr a b) x -> ZippedStackRepr a b)
-> Generic (ZippedStackRepr a b)
forall x. Rep (ZippedStackRepr a b) x -> ZippedStackRepr a b
forall x. ZippedStackRepr a b -> Rep (ZippedStackRepr a b) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a b x. Rep (ZippedStackRepr a b) x -> ZippedStackRepr a b
forall a b x. ZippedStackRepr a b -> Rep (ZippedStackRepr a b) x
$cto :: forall a b x. Rep (ZippedStackRepr a b) x -> ZippedStackRepr a b
$cfrom :: forall a b x. ZippedStackRepr a b -> Rep (ZippedStackRepr a b) x
Generic)
  deriving anyclass (WellTypedToT (ZippedStackRepr a b)
WellTypedToT (ZippedStackRepr a b)
-> (ZippedStackRepr a b -> Value (ToT (ZippedStackRepr a b)))
-> (Value (ToT (ZippedStackRepr a b)) -> ZippedStackRepr a b)
-> IsoValue (ZippedStackRepr a b)
Value (ToT (ZippedStackRepr a b)) -> ZippedStackRepr a b
ZippedStackRepr a b -> Value (ToT (ZippedStackRepr a b))
forall a.
WellTypedToT a
-> (a -> Value (ToT a)) -> (Value (ToT a) -> a) -> IsoValue a
forall {a} {b}.
(IsoValue a, IsoValue b) =>
WellTypedToT (ZippedStackRepr a b)
forall a b.
(IsoValue a, IsoValue b) =>
Value (ToT (ZippedStackRepr a b)) -> ZippedStackRepr a b
forall a b.
(IsoValue a, IsoValue b) =>
ZippedStackRepr a b -> Value (ToT (ZippedStackRepr a b))
fromVal :: Value (ToT (ZippedStackRepr a b)) -> ZippedStackRepr a b
$cfromVal :: forall a b.
(IsoValue a, IsoValue b) =>
Value (ToT (ZippedStackRepr a b)) -> ZippedStackRepr a b
toVal :: ZippedStackRepr a b -> Value (ToT (ZippedStackRepr a b))
$ctoVal :: forall a b.
(IsoValue a, IsoValue b) =>
ZippedStackRepr a b -> Value (ToT (ZippedStackRepr a b))
IsoValue, AnnOptions
FollowEntrypointFlag -> Notes (ToT (ZippedStackRepr a b))
(FollowEntrypointFlag -> Notes (ToT (ZippedStackRepr a b)))
-> AnnOptions -> HasAnnotation (ZippedStackRepr a b)
forall a.
(FollowEntrypointFlag -> Notes (ToT a))
-> AnnOptions -> HasAnnotation a
forall a b. (HasAnnotation a, HasAnnotation b) => AnnOptions
forall a b.
(HasAnnotation a, HasAnnotation b) =>
FollowEntrypointFlag -> Notes (ToT (ZippedStackRepr a b))
annOptions :: AnnOptions
$cannOptions :: forall a b. (HasAnnotation a, HasAnnotation b) => AnnOptions
getAnnotation :: FollowEntrypointFlag -> Notes (ToT (ZippedStackRepr a b))
$cgetAnnotation :: forall a b.
(HasAnnotation a, HasAnnotation b) =>
FollowEntrypointFlag -> Notes (ToT (ZippedStackRepr a b))
HasAnnotation)
-- | A type used to represent an empty zipped stack, isomorphic to a unit and
-- represented as such in Michelson.
data ZSNil = ZSNil
  deriving stock (Int -> ZSNil -> ShowS
[ZSNil] -> ShowS
ZSNil -> String
(Int -> ZSNil -> ShowS)
-> (ZSNil -> String) -> ([ZSNil] -> ShowS) -> Show ZSNil
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ZSNil] -> ShowS
$cshowList :: [ZSNil] -> ShowS
show :: ZSNil -> String
$cshow :: ZSNil -> String
showsPrec :: Int -> ZSNil -> ShowS
$cshowsPrec :: Int -> ZSNil -> ShowS
Show, ZSNil -> ZSNil -> Bool
(ZSNil -> ZSNil -> Bool) -> (ZSNil -> ZSNil -> Bool) -> Eq ZSNil
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ZSNil -> ZSNil -> Bool
$c/= :: ZSNil -> ZSNil -> Bool
== :: ZSNil -> ZSNil -> Bool
$c== :: ZSNil -> ZSNil -> Bool
Eq, (forall x. ZSNil -> Rep ZSNil x)
-> (forall x. Rep ZSNil x -> ZSNil) -> Generic ZSNil
forall x. Rep ZSNil x -> ZSNil
forall x. ZSNil -> Rep ZSNil x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep ZSNil x -> ZSNil
$cfrom :: forall x. ZSNil -> Rep ZSNil x
Generic)
  deriving anyclass (WellTypedToT ZSNil
WellTypedToT ZSNil
-> (ZSNil -> Value (ToT ZSNil))
-> (Value (ToT ZSNil) -> ZSNil)
-> IsoValue ZSNil
Value (ToT ZSNil) -> ZSNil
ZSNil -> Value (ToT ZSNil)
forall a.
WellTypedToT a
-> (a -> Value (ToT a)) -> (Value (ToT a) -> a) -> IsoValue a
fromVal :: Value (ToT ZSNil) -> ZSNil
$cfromVal :: Value (ToT ZSNil) -> ZSNil
toVal :: ZSNil -> Value (ToT ZSNil)
$ctoVal :: ZSNil -> Value (ToT ZSNil)
IsoValue, AnnOptions
FollowEntrypointFlag -> Notes (ToT ZSNil)
(FollowEntrypointFlag -> Notes (ToT ZSNil))
-> AnnOptions -> HasAnnotation ZSNil
forall a.
(FollowEntrypointFlag -> Notes (ToT a))
-> AnnOptions -> HasAnnotation a
annOptions :: AnnOptions
$cannOptions :: AnnOptions
getAnnotation :: FollowEntrypointFlag -> Notes (ToT ZSNil)
$cgetAnnotation :: FollowEntrypointFlag -> Notes (ToT ZSNil)
HasAnnotation)

instance (Buildable a, Buildable b) => Buildable (ZippedStackRepr a b) where
  build :: ZippedStackRepr a b -> Builder
build (a
a ::: b
b) = a
a a -> Builder -> Builder
forall a b. (Buildable a, FromBuilder b) => a -> Builder -> b
|+ Builder
" : " Builder -> Builder -> Builder
forall b. FromBuilder b => Builder -> Builder -> b
+| b
b b -> Builder -> Builder
forall a b. (Buildable a, FromBuilder b) => a -> Builder -> b
|+ Builder
""

instance Buildable ZSNil where
  build :: ZSNil -> Builder
build ZSNil
ZSNil = Builder
"[]"

-- | Zipping stack into tuple and back.
class (KnownIsoT (ZippedStack s)) => ZipInstr (s :: [Type]) where
  -- | A type which contains the whole stack zipped.
  type ZippedStack s :: 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 :: forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr = Instr (ToTs s) (ToTs '[ZippedStack s]) -> s :-> '[ZippedStack s]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (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 :: forall (s :: [*]). ZipInstr s => '[ZippedStack s] :-> s
unzipInstr = Instr (ToTs '[ZippedStack s]) (ToTs s) -> '[ZippedStack s] :-> s
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (forall (s :: [*]).
ZipInstr s =>
Instr '[ToT (ZippedStack s)] (ToTs s)
unzipInstrTyped @s)

instance ZipInstr '[] where
  type ZippedStack '[] = ZSNil
  zipInstrTyped :: Instr (ToTs '[]) '[ToT (ZippedStack '[])]
zipInstrTyped = Instr (ToTs '[]) '[ToT (ZippedStack '[])]
forall {inp :: [T]} {out :: [T]} (s :: [T]).
(inp ~ s, out ~ ('TUnit : s)) =>
Instr inp out
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) = ZippedStackRepr 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 (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 {inp :: [T]} {out :: [T]} (a :: T) (b :: T) (s :: [T]).
(inp ~ (a : b : s), out ~ ('TPair a b : s)) =>
Instr inp out
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 {inp :: [T]} {out :: [T]} (a :: T) (b :: T) (s :: [T]).
(inp ~ ('TPair a b : s), out ~ (a : b : s)) =>
Instr inp out
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 (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 -> Fn (ZippedStack inp) (ZippedStack out)
zippingStack :: forall (inp :: [*]) (out :: [*]).
ZipInstrs '[inp, out] =>
(inp :-> out) -> Fn (ZippedStack inp) (ZippedStack out)
zippingStack 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])
-> '[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]
  => Fn (ZippedStack inp) (ZippedStack out) -> inp :-> out
unzippingStack :: forall (inp :: [*]) (out :: [*]).
ZipInstrs '[inp, out] =>
Fn (ZippedStack inp) (ZippedStack out) -> inp :-> out
unzippingStack Fn (ZippedStack inp) (ZippedStack out)
code = inp :-> '[ZippedStack inp]
forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr (inp :-> '[ZippedStack inp])
-> Fn (ZippedStack inp) (ZippedStack out)
-> inp :-> '[ZippedStack out]
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## Fn (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