-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

{- | Reimplementation of some syntax sugar.

You need the following module pragmas to make it work smoothly:

{-# LANGUAGE NoApplicativeDo, RebindableSyntax #-}
{-# OPTIONS_GHC -Wno-unused-do-bind #-}

-}
module Lorentz.Rebinded
  ( (>>)
  , pure
  , return
  , ifThenElse
  , Condition (..)
  , (<.)
  , (>.)
  , (<=.)
  , (>=.)
  , (==.)
  , (/=.)
  , keepIfArgs

    -- * Re-exports required for RebindableSyntax
  , fromInteger
  , fromString
  , fromLabel
  ) where


import Prelude hiding (drop, swap, (>>), (>>=))

import Named ((:!))

import Lorentz.Arith
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Instr
import Lorentz.Macro
import Michelson.Typed.Arith
import Util.Label (Label)

-- | Aliases for '(#)' used by do-blocks.
(>>) :: (a :-> b) -> (b :-> c) -> (a :-> c)
(>>) = (#)

-- | Predicate for @if ... then .. else ...@ construction,
-- defines a kind of operation applied to the top elements of the current stack.
--
-- Type arguments mean:
-- 1. Input of @if@
-- 2. Left branch input
-- 3. Right branch input
-- 4. Output of branches
-- 5. Output of @if@
data Condition arg argl argr outb out where
  Holds :: Condition (Bool ': s) s s o o
  IsSome :: Condition (Maybe a ': s) (a ': s) s o o
  IsNone :: Condition (Maybe a ': s) s (a ': s) o o
  IsLeft :: Condition (Either l r ': s) (l ': s) (r ': s) o o
  IsRight :: Condition (Either l r ': s) (r ': s) (l ': s) o o
  IsCons :: Condition ([a] ': s) (a ': [a] ': s) s o o
  IsNil :: Condition ([a] ': s) s (a ': [a] ': s) o o

  IsZero :: (UnaryArithOpHs Eq' a, UnaryArithResHs Eq' a ~ Bool)
         => Condition (a ': s) s s o o
  IsNotZero :: (UnaryArithOpHs Eq' a, UnaryArithResHs Eq' a ~ Bool)
         => Condition (a ': s) s s o o

  IsEq :: NiceComparable a => Condition (a ': a ': s) s s o o
  IsNeq :: NiceComparable a => Condition (a ': a ': s) s s o o
  IsLt :: NiceComparable a => Condition (a ': a ': s) s s o o
  IsGt :: NiceComparable a => Condition (a ': a ': s) s s o o
  IsLe :: NiceComparable a => Condition (a ': a ': s) s s o o
  IsGe :: NiceComparable a => Condition (a ': a ': s) s s o o

  -- | Explicitly named binary condition, to ensure proper order of
  -- stack arguments.
  NamedBinCondition ::
    Condition (a ': a ': s) s s o o ->
    Label n1 -> Label n2 ->
    Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o

  -- | Provide the compared arguments to @if@ branches.
  PreserveArgsBinCondition ::
    (forall st o. Condition (a ': b ': st) st st o o) ->
    Condition (a ': b ': s) (a ': b ': s) (a ': b ': s) (a ': b ': s) s

-- | Defines semantics of @if ... then ... else ...@ construction.
ifThenElse
  :: Condition arg argl argr outb out
  -> (argl :-> outb) -> (argr :-> outb) -> (arg :-> out)
ifThenElse = \case
  Holds -> if_
  IsSome -> flip ifNone
  IsNone -> ifNone
  IsLeft -> ifLeft
  IsRight -> flip ifLeft
  IsCons -> ifCons
  IsNil -> flip ifCons

  IsZero -> \l r -> eq0 # if_ l r
  IsNotZero -> \l r -> eq0 # if_ r l

  IsEq -> ifEq
  IsNeq -> ifNeq
  IsLt -> ifLt
  IsGt -> ifGt
  IsLe -> ifLe
  IsGe -> ifGe

  NamedBinCondition condition l1 l2 -> \l r ->
    fromNamed l1 # dip (fromNamed l2) # ifThenElse condition l r

  PreserveArgsBinCondition condition -> \l r ->
    dip dup # swap # dip dup # swap #
    ifThenElse condition
      -- since this pattern is commonly used when one of the branches fails,
      -- it's essential to @drop@ within branches, not after @if@ - @drop@s
      -- appearing to be dead code will be cut off
      (l # drop # drop)
      (r # drop # drop)

-- | Named version of 'IsLt'.
--
-- In this and similar operators you provide names of accepted stack operands as
-- a safety measure of that they go in the expected order.
infix 4 <.
(<.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(<.) = NamedBinCondition IsLt

-- | Named version of 'IsGt'.
infix 4 >.
(>.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(>.) = NamedBinCondition IsGt

-- | Named version of 'IsLe'.
infix 4 <=.
(<=.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(<=.) = NamedBinCondition IsLe

-- | Named version of 'IsGe'.
infix 4 >=.
(>=.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(>=.) = NamedBinCondition IsGe

-- | Named version of 'IsEq'.
infix 4 ==.
(==.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(==.) = NamedBinCondition IsEq

-- | Named version of 'IsNeq'.
infix 4 /=.
(/=.)
  :: NiceComparable a
  => Label n1 -> Label n2
  -> Condition ((n1 :! a) ': (n2 :! a) ': s) s s o o
(/=.) = NamedBinCondition IsNeq

-- | Condition modifier, makes stack operands of binary comparison to be
-- available within @if@ branches.
keepIfArgs
  :: (forall st o. Condition (a ': b ': st) st st o o)
  -> Condition (a ': b ': s) (a ': b ': s) (a ': b ': s) (a ': b ': s) s
keepIfArgs = PreserveArgsBinCondition