{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Sum where

import Data.Vector.Fixed (S,Z,Fun(..))
import Data.Vector.Fixed.Cont (Fn,Arity(..))

-- | Type family for sum of unary natural numbers.
type family n + m :: *

type instance Z   + n = n
type instance S n + k = S (n + k)


-- | We need to extend Arity type class with ability generate proofs
--   that @Fn (n+k) a b ~ Fn n a (Fn k a b)@ which is needed for
--   'uncurryMany'.
class Arity n => AritySum n where
  witSum :: WitSum n k a b

instance AritySum Z where
  witSum = WitSum
instance AritySum n => AritySum (S n) where
  witSum :: forall k a b. WitSum (S n) k a b
  witSum = case witSum :: WitSum n k a b of
             WitSum -> WitSum

-- | Value that carry proof that `Fn (n+k) a b ~ Fn n a (Fn k a b)`
data WitSum n k a b where
  WitSum :: (Fn (n+k) a b ~ Fn n a (Fn k a b)) => WitSum n k a b


uncurryMany :: forall n k a b. AritySum n
            => Fun n a (Fun k a b) -> Fun (n + k) a b
uncurryMany f =
  case witSum :: WitSum n k a b of
    WitSum ->
      case fmap unFun f :: Fun n a (Fn k a b) of
        Fun g -> Fun g


curryMany :: forall n k a b. Arity n
          => Fun (n + k) a b -> Fun n a (Fun k a b)
curryMany (Fun f0) = Fun $ accum
  (\(T_curry f) a -> T_curry (f a))
  (\(T_curry f)   -> Fun f :: Fun k a b)
  (  T_curry f0 :: T_curry a b k n)

newtype T_curry a b k n = T_curry (Fn (n + k) a b)