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

{- | 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
  , IsCondition (ifThenElse)
  , Condition (..)
  , (<.)
  , (>.)
  , (<=.)
  , (>=.)
  , (==.)
  , (/=.)
  , keepIfArgs

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


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

import Lorentz.Arith
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Constraints.Scopes
import Lorentz.Instr
import Lorentz.Macro
import Morley.Michelson.Typed.Arith
import Morley.Util.Label (Label)
import Morley.Util.Named

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

-- | The most basic 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

  Not :: Condition s s1 s2 ob o -> Condition s s2 s1 ob o

  IsZero :: (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 ::
    (Dupable a, Dupable b) =>
    (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

-- | Everything that can be put after @if@ keyword.
--
-- The first type argument stands for the condition type, and all other type
-- arguments define stack types around/within the @if then else@ construction.
-- For semantics of each type argument see 'Condition'.
class IsCondition cond arg argl argr outb out where
  -- | Defines semantics of @if ... then ... else ...@ construction.
  ifThenElse :: cond -> (argl :-> outb) -> (argr :-> outb) -> (arg :-> out)

instance (arg ~ arg0, argl ~ argl0, argr ~ argr0, outb ~ outb0, out ~ out0) =>
         IsCondition (Condition arg argl argr outb out) arg0 argl0 argr0 outb0 out0 where
  ifThenElse = \case
    Holds -> if_
    IsSome -> flip ifNone
    IsNone -> ifNone
    IsLeft -> ifLeft
    IsRight -> flip ifLeft
    IsCons -> ifCons
    IsNil -> flip ifCons

    Not cond -> \l r -> ifThenElse cond r l

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

    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 ->
      dupN @2 # dupN @2 #
      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
  :: (Dupable a, Dupable b)
  => (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