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

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

import Prelude hiding (drop)

import Lorentz.Annotation
import Lorentz.Base
import Morley.AsRPC (HasRPCRepr(..))
import Morley.Michelson.Typed
import Morley.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 ## I r = I (seqOpt l r)
l ## r = l # r

seqOpt :: Instr a b -> Instr b c -> Instr a c
seqOpt l r =
  -- We are very verbose about cases to avoid
  -- significant compilation time increase
  case l of
    Nop -> r
    DIP Nop -> r
    x -> case r of
      Nop -> x
      DIP Nop -> x
      _ -> l `Seq` 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 (Show, Eq, Generic)
  deriving anyclass (IsoValue, 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 (Show, Eq, Generic)
  deriving anyclass (IsoValue, HasAnnotation)

-- | 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 = I (zipInstrTyped @s)

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

instance ZipInstr '[] where
  type ZippedStack '[] = ZSNil
  zipInstrTyped = UNIT
  unzipInstrTyped = DROP

instance (KnownIsoT a) => ZipInstr '[a] where
  type ZippedStack '[a] = a
  zipInstrTyped = Nop
  unzipInstrTyped = 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 = DIP (zipInstrTyped @(b ': s)) `seqOpt` PAIR
  unzipInstrTyped = UNPAIR `seqOpt` DIP (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 code = unzipInstr ## code ## zipInstr

-- | Unflatten both ends of instruction stack.
unzippingStack
  :: ZipInstrs [inp, out]
  => Fn (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 = mkVLam $ 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)

instance HasRPCRepr (inp :-> out) where
  type AsRPC (inp :-> out) = inp :-> out