{-# 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.Base
import Michelson.Typed

-- | Zipping stack into tuple and back.
class 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.
-}

-- | Version of '#' which performs some optimizations immediatelly.
-- Need it here at least because we don't want to surround
-- @fromVal (Lambda i o)@ with nops.
(##) :: (a :-> b) -> (b :-> c) -> (a :-> c)
l ## r = case (l, r) of
  (I Nop, I x) -> I x
  (I x, I Nop) -> I x
  _ -> l # r

instance ZipInstr '[] where
  type ZippedStack '[] = ()
  zipInstr = I UNIT
  unzipInstr = I DROP

instance ZipInstr '[a] where
  type ZippedStack '[a] = a
  zipInstr = I Nop
  unzipInstr = I 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 = I PAIR ## zipInstr @((a, b) ': s)
  unzipInstr = unzipInstr @((a, b) ': s) ## I (DUP `Seq` CAR `Seq` DIP 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 code = unzipInstr ## code ## zipInstr

-- | Unflatten both ends of instruction stack.
unzippingStack
  :: ZipInstrs [inp, out]
  => Lambda (ZippedStack inp) (ZippedStack out) -> inp :-> out
unzippingStack code = zipInstr ## code ## unzipInstr

instance (ZipInstr inp, ZipInstr out) => IsoValue (inp :-> out) where
  type ToT (inp :-> out) = 'TLambda (ToT (ZippedStack inp)) (ToT (ZippedStack out))
  toVal i = VLam . unLorentzInstr $ zippingStack i
  fromVal (VLam i) = zipInstr ## LorentzInstr i ## unzipInstr