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

-- | Constraints for some typed instructions we share accross a few places and
-- related type-level helpers.
module Morley.Michelson.Typed.Instr.Constraints
  ( ConstraintDUPN'
  , ConstraintDUPN
  , ConstraintDIPN'
  , ConstraintDIPN
  , ConstraintDIG'
  , ConstraintDIG
  , ConstraintDUG'
  , ConstraintDUG
  , ConstraintPairN
  , ConstraintUnpairN
  , ConstraintGetN
  , ConstraintUpdateN

    -- * Helpers
  , PairN
  , UnpairN
  , UpdateN
  , GetN
  , RightComb
  ) where

import Prelude hiding (EQ, GT, LT)

import GHC.TypeNats (type (+))

import Morley.Michelson.Typed.T (T(..))
import Morley.Michelson.Typed.TypeLevel
  (CombedPairLeafCount, CombedPairLeafCountIsAtLeast, CombedPairNodeCount,
  CombedPairNodeIndexIsValid, IsPair)
import Morley.Util.Peano
import Morley.Util.Type (FailUnless, If, type (++))
import Morley.Util.TypeLits (ErrorMessage(ShowType, Text, (:$$:), (:<>:)))

-- | Constraint that is used in DUPN, we want to share it with
-- typechecking code and eDSL code.
type ConstraintDUPN' kind (n :: Peano) (inp :: [kind]) (out :: [kind]) (a :: kind) =
  ( RequireLongerOrSameLength inp n, n > 'Z ~ 'True
  , inp ~ (LazyTake (Decrement n) inp ++ (a ': Drop n inp))
  , out ~ (a ': inp)
  )

type ConstraintDUPN n inp out a = ConstraintDUPN' T n inp out a

-- | Constraint that is used in DIPN, we want to share it with
-- typechecking code and eDSL code.
type ConstraintDIPN' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (s :: [kind]) (s' :: [kind]) =
  ( RequireLongerOrSameLength inp n
  , LazyTake n inp ++ s ~ inp
  , LazyTake n inp ++ s' ~ out
  , s ~ Drop n inp
  , s' ~ Drop n out
  )

type ConstraintDIPN n inp out s s' = ConstraintDIPN' T n inp out s s'

-- | Constraint that is used in DIG, we want to share it with
-- typechecking code and eDSL code.
type ConstraintDIG' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (a :: kind) =
  ( RequireLongerThan inp n
  , out ~ a ': LazyTake n inp ++ Drop ('S n) inp
  , inp ~ LazyTake n (Drop ('S 'Z) out) ++ (a ': Drop ('S n) out)
  )

type ConstraintDIG n inp out a = ConstraintDIG' T n inp out a

-- | Constraint that is used in DUG, we want to share it with
-- typechecking code and eDSL code.
type ConstraintDUG' kind (n :: Peano) (inp :: [kind])
  (out :: [kind]) (a :: kind) =
  ConstraintDIG' kind n out inp a

type ConstraintDUG n inp out a = ConstraintDUG' T n inp out a

type ConstraintPairN (n :: Peano) (inp :: [T]) =
  ( RequireLongerOrSameLength inp n
  , FailUnless (n >= ToPeano 2) ('Text "'PAIR n' expects n ≥ 2")
  )

type PairN (n :: Peano) (s :: [T]) = RightComb (LazyTake n s) ': Drop n s

-- | Folds a stack into a right-combed pair.
--
-- > RightComb '[ 'TInt,  'TString, 'TUnit ]
-- > ~
-- > 'TPair 'TInt ('TPair 'TString 'TUnit)
type family RightComb (s :: [T]) :: T where
  RightComb '[ x, y ] = 'TPair x y
  RightComb (x ': xs) = 'TPair x (RightComb xs)

type ConstraintUnpairN (n :: Peano) (pair :: T) =
  ( FailUnless (n >= ToPeano 2)
      ('Text "'UNPAIR n' expects n ≥ 2")

  , FailUnless (CombedPairLeafCountIsAtLeast n pair)
      (If (IsPair pair)
        ('Text "'UNPAIR "
          ':<>: 'ShowType (FromPeano n)
          ':<>: 'Text "' expects a right-combed pair with at least "
          ':<>: 'ShowType (FromPeano n)
          ':<>: 'Text " elements at the top of the stack,"
          ':$$: 'Text "but the pair only contains "
          ':<>: 'ShowType (FromPeano (CombedPairLeafCount pair))
          ':<>: 'Text " elements.")
        ('Text "Expected a pair at the top of the stack, but found: "
          ':<>: 'ShowType pair
        )
      )
  )

-- | Splits a right-combed pair into @n@ elements.
--
-- > UnpairN (ToPeano 3) ('TPair 'TInt ('TPair 'TString 'TUnit))
-- > ~
-- > '[ 'TInt, 'TString, 'TUnit]
--
-- > UnpairN (ToPeano 3) ('TPair 'TInt ('TPair 'TString ('TPair 'TUnit 'TInt)))
-- > ~
-- > '[ 'TInt, 'TString, 'TPair 'TUnit 'TInt]
type family UnpairN (n :: Peano) (s :: T) :: [T] where
  UnpairN ('S ('S 'Z)) ('TPair x y) = [x, y]
  UnpairN ('S n)       ('TPair x y) = x : UnpairN n y

type ConstraintGetN (ix :: Peano) (pair :: T) =
  ( FailUnless (CombedPairNodeIndexIsValid ix pair)
      (If (IsPair pair)
        ('Text "'GET "
          ':<>: 'ShowType (FromPeano ix)
          ':<>: 'Text "' expects a right-combed pair with at least "
          ':<>: 'ShowType (FromPeano ix + 1)
          ':<>: 'Text " nodes at the top of the stack,"
          ':$$: 'Text "but the pair only contains "
          ':<>: 'ShowType (FromPeano (CombedPairNodeCount pair))
          ':<>: 'Text " nodes.")
        ('Text "Expected a pair at the top of the stack, but found: "
          ':<>: 'ShowType pair
        )
      )
  )

-- | Get the node at index @ix@ of a right-combed pair.
type family GetN (ix :: Peano) (pair :: T) :: T where
  GetN 'Z val                       = val
  GetN ('S 'Z) ('TPair left _)      = left
  GetN ('S ('S n)) ('TPair _ right) = GetN n right

type ConstraintUpdateN (ix :: Peano) (pair :: T) =
  ( FailUnless (CombedPairNodeIndexIsValid ix pair)
      (If (IsPair pair)
        ('Text "'UPDATE "
          ':<>: 'ShowType (FromPeano ix)
          ':<>: 'Text "' expects the 2nd element of the stack to be a right-combed pair with at least "
          ':<>: 'ShowType (FromPeano ix + 1)
          ':<>: 'Text " nodes,"
          ':$$: 'Text "but the pair only contains "
          ':<>: 'ShowType (FromPeano (CombedPairNodeCount pair))
          ':<>: 'Text " nodes.")
        ('Text "Expected the 2nd element of the stack to be a pair, but found: "
          ':<>: 'ShowType pair
        )
      )
  )

-- | Update the node at index @ix@ of a right-combed pair.
type family UpdateN (ix :: Peano) (val :: T) (pair :: T) :: T where
  UpdateN 'Z          val _                   = val
  UpdateN ('S 'Z)     val ('TPair _  right)   = 'TPair val right
  UpdateN ('S ('S n)) val ('TPair left right) = 'TPair left (UpdateN n val right)