-- 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 = I UNIT
  unzipInstr = I DROP

instance (KnownIsoT a) => 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 (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 i = VLam . unLorentzInstr $ zippingStack i
  fromVal (VLam i) = zipInstr ## LorentzInstr i ## unzipInstr

instance
    ( HasAnnotation (ZippedStack i)
    , HasAnnotation (ZippedStack o)
    )
  =>
    HasAnnotation (i :-> o)
  where
  getAnnotation b = NTLambda noAnn
    (getAnnotation @(ZippedStack i) b)
    (getAnnotation @(ZippedStack o) b)