-- 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 (..)
  , 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)

-- | 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.
  zipInstr :: s :-> '[ZippedStack s]

  -- | Unfold given stack from a single value.
  unzipInstr :: '[ZippedStack s] :-> s

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

instance ZipInstr '[] where
  type ZippedStack '[] = ()
  zipInstr :: '[] :-> '[ZippedStack '[]]
zipInstr = Instr (ToTs '[]) (ToTs '[()]) -> '[] :-> '[()]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I Instr (ToTs '[]) (ToTs '[()])
forall (inp :: [T]). Instr inp ('TUnit : inp)
UNIT
  unzipInstr :: '[ZippedStack '[]] :-> '[]
unzipInstr = Instr (ToTs '[()]) (ToTs '[]) -> '[()] :-> '[]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I Instr (ToTs '[()]) (ToTs '[])
forall (a :: T) (out :: [T]). Instr (a : out) out
DROP

instance (KnownIsoT a) => ZipInstr '[a] where
  type ZippedStack '[a] = a
  zipInstr :: '[a] :-> '[ZippedStack '[a]]
zipInstr = Instr (ToTs '[a]) (ToTs '[a]) -> '[a] :-> '[a]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I Instr (ToTs '[a]) (ToTs '[a])
forall (inp :: [T]). Instr inp inp
Nop
  unzipInstr :: '[ZippedStack '[a]] :-> '[a]
unzipInstr = Instr (ToTs '[a]) (ToTs '[a]) -> '[a] :-> '[a]
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I Instr (ToTs '[a]) (ToTs '[a])
forall (inp :: [T]). Instr inp inp
Nop

-- | Such definition seems the only possible one we can support
-- efficiently.
instance ZipInstr ((a, b) ': s) => ZipInstr (a ': b ': s) where
  type ZippedStack (a ': b ': s) = ZippedStack ((a, b) ': s)
  zipInstr :: (a : b : s) :-> '[ZippedStack (a : b : s)]
zipInstr = Instr (ToTs (a : b : s)) (ToTs ((a, b) : s))
-> (a : b : s) :-> ((a, b) : s)
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I Instr (ToTs (a : b : s)) (ToTs ((a, 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 ((a : b : s) :-> ((a, b) : s))
-> (((a, b) : s) :-> '[ZippedStack ((a, b) : s)])
-> (a : b : s) :-> '[ZippedStack ((a, b) : s)]
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## ZipInstr ((a, b) : s) =>
((a, b) : s) :-> '[ZippedStack ((a, b) : s)]
forall (s :: [*]). ZipInstr s => s :-> '[ZippedStack s]
zipInstr @((a, b) ': s)
  unzipInstr :: '[ZippedStack (a : b : s)] :-> (a : b : s)
unzipInstr = ZipInstr ((a, b) : s) =>
'[ZippedStack ((a, b) : s)] :-> ((a, b) : s)
forall (s :: [*]). ZipInstr s => '[ZippedStack s] :-> s
unzipInstr @((a, b) ': s) ('[ZippedStack ((a, b) : s)] :-> ((a, b) : s))
-> (((a, b) : s) :-> (a : b : s))
-> '[ZippedStack ((a, b) : s)] :-> (a : b : s)
forall (a :: [*]) (b :: [*]) (c :: [*]).
(a :-> b) -> (b :-> c) -> a :-> c
## Instr (ToTs ((a, b) : s)) (ToTs (a : b : s))
-> ((a, b) : s) :-> (a : b : s)
forall (inp :: [*]) (out :: [*]).
Instr (ToTs inp) (ToTs out) -> inp :-> out
I (Instr
  ('TPair (ToT a) (ToT b) : ToTs s)
  ('TPair (ToT a) (ToT b) : 'TPair (ToT a) (ToT b) : ToTs s)
forall (a :: T) (s :: [T]). Instr (a : s) (a : a : s)
DUP Instr
  ('TPair (ToT a) (ToT b) : ToTs s)
  ('TPair (ToT a) (ToT b) : 'TPair (ToT a) (ToT b) : ToTs s)
-> Instr
     ('TPair (ToT a) (ToT b) : 'TPair (ToT a) (ToT b) : ToTs s)
     (ToT a : 'TPair (ToT a) (ToT b) : ToTs s)
-> Instr
     ('TPair (ToT a) (ToT b) : ToTs s)
     (ToT a : 'TPair (ToT a) (ToT b) : ToTs s)
forall (inp :: [T]) (b :: [T]) (out :: [T]).
Instr inp b -> Instr b out -> Instr inp out
`Seq` Instr
  ('TPair (ToT a) (ToT b) : 'TPair (ToT a) (ToT b) : ToTs s)
  (ToT a : 'TPair (ToT a) (ToT b) : ToTs s)
forall (i :: [T]) (o :: [T]) (a :: T) (b :: T) (s :: [T]).
(i ~ ('TPair a b : s), o ~ (a : s)) =>
Instr i o
CAR Instr
  ('TPair (ToT a) (ToT b) : ToTs s)
  (ToT a : 'TPair (ToT a) (ToT b) : ToTs s)
-> Instr
     (ToT a : 'TPair (ToT a) (ToT b) : ToTs s) (ToT a : ToT b : ToTs s)
-> Instr ('TPair (ToT a) (ToT b) : ToTs s) (ToT a : ToT b : ToTs s)
forall (inp :: [T]) (b :: [T]) (out :: [T]).
Instr inp b -> Instr b out -> Instr inp out
`Seq` Instr ('TPair (ToT a) (ToT b) : ToTs s) (ToT b : ToTs s)
-> Instr
     (ToT a : 'TPair (ToT a) (ToT b) : ToTs s) (ToT a : ToT b : ToTs s)
forall (a :: [T]) (c :: [T]) (b :: T).
Instr a c -> Instr (b : a) (b : c)
DIP Instr ('TPair (ToT a) (ToT b) : ToTs s) (ToT b : ToTs s)
forall (i :: [T]) (o :: [T]) (a :: T) (b :: T) (s :: [T]).
(i ~ ('TPair a b : s), o ~ (b : s)) =>
Instr i o
CDR)

-- | 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)