{-# LANGUAGE UndecidableInstances #-}

-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at https://mozilla.org/MPL/2.0/.

{- |
Copyright   :  (c) 2023 Yamada Ryo
License     :  MPL-2.0 (see the file LICENSE)
Maintainer  :  ymdfield@outlook.jp
Stability   :  experimental
Portability :  portable

An implementation of an open union for first-order effects using
the [extensible](https://hackage.haskell.org/package/extensible) package as a backend.
-}
module Data.Free.Extensible where

import Control.Effect.Class (Instruction)
import Data.Extensible (Forall, Match (Match), htabulateFor, leadership, match)
import Data.Extensible.Sum (exhaust, strikeAt, (<:|), type (:/) (EmbedAt))
import Data.Free.Union (
    Union (
        HasMembership,
        absurdUnion,
        inject,
        inject0,
        project,
        weaken,
        (|+|:)
    ),
 )
import Data.Proxy (Proxy (Proxy))
import GHC.TypeNats (KnownNat, Nat, natVal, type (+))
import Type.Membership (Membership, nextMembership)
import Unsafe.Coerce (unsafeCoerce)

{- |
An implementation of an open union for first-order effects using
the [extensible](https://hackage.haskell.org/package/extensible) package as a backend.
-}
newtype ExtensibleUnion fs a = ExtensibleUnion {forall (fs :: [* -> *]) a. ExtensibleUnion fs a -> fs :/ FieldApp a
unExtensibleUnion :: fs :/ FieldApp a}

newtype FieldApp a (f :: Instruction) = FieldApp {forall a (f :: * -> *). FieldApp a f -> f a
unFieldApp :: f a}

instance Forall Functor fs => Functor (ExtensibleUnion fs) where
    fmap :: forall a b.
(a -> b) -> ExtensibleUnion fs a -> ExtensibleUnion fs b
fmap a -> b
f =
        forall (fs :: [* -> *]) a.
(fs :/ FieldApp a) -> ExtensibleUnion fs a
ExtensibleUnion
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (xs :: [k]) (h :: k -> *) a.
(xs :& Match h a) -> (xs :/ h) -> a
match
                ( forall {k} (c :: k -> Constraint) (xs :: [k])
       (proxy :: (k -> Constraint) -> *) (h :: k -> *).
Forall c xs =>
proxy c
-> (forall (x :: k). c x => Membership xs x -> h x) -> xs :& h
htabulateFor @Functor forall {k} (t :: k). Proxy t
Proxy \Membership fs x
w ->
                    forall {k} (h :: k -> *) r (x :: k). (h x -> r) -> Match h r x
Match \FieldApp a x
e -> forall {k} (xs :: [k]) (x :: k) (h :: k -> *).
Membership xs x -> h x -> xs :/ h
EmbedAt Membership fs x
w forall a b. (a -> b) -> a -> b
$ forall a (f :: * -> *). f a -> FieldApp a f
FieldApp forall a b. (a -> b) -> a -> b
$ a -> b
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a (f :: * -> *). FieldApp a f -> f a
unFieldApp FieldApp a x
e
                )
            forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (fs :: [* -> *]) a. ExtensibleUnion fs a -> fs :/ FieldApp a
unExtensibleUnion
    {-# INLINE fmap #-}

{- todo:
instance Forall Foldable fs => Foldable (ExtensibleUnion fs) where
instance Forall Traversable fs => Traversable (ExtensibleUnion fs) where
-}

instance Union ExtensibleUnion where
    type HasMembership _ f fs = KnownNat (TypeIndex fs f)

    inject :: forall (f :: * -> *) (fs :: [* -> *]).
HasMembership ExtensibleUnion f fs =>
f ~> ExtensibleUnion fs
inject = forall (fs :: [* -> *]) a.
(fs :/ FieldApp a) -> ExtensibleUnion fs a
ExtensibleUnion forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (xs :: [k]) (x :: k) (h :: k -> *).
Membership xs x -> h x -> xs :/ h
EmbedAt forall {k} (xs :: [k]) (x :: k).
KnownNat (TypeIndex xs x) =>
Membership xs x
findFirstMembership forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (f :: * -> *). f a -> FieldApp a f
FieldApp
    {-# INLINE inject #-}

    project :: forall (f :: * -> *) (fs :: [* -> *]) a.
HasMembership ExtensibleUnion f fs =>
ExtensibleUnion fs a -> Maybe (f a)
project (ExtensibleUnion fs :/ FieldApp a
u) = forall a (f :: * -> *). FieldApp a f -> f a
unFieldApp forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall {k} (h :: k -> *) (x :: k) (xs :: [k]).
Membership xs x -> (xs :/ h) -> Maybe (h x)
strikeAt forall {k} (xs :: [k]) (x :: k).
KnownNat (TypeIndex xs x) =>
Membership xs x
findFirstMembership fs :/ FieldApp a
u
    {-# INLINE project #-}

    absurdUnion :: forall a x. ExtensibleUnion '[] a -> x
absurdUnion = forall {k} (h :: k -> *) r. ('[] :/ h) -> r
exhaust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (fs :: [* -> *]) a. ExtensibleUnion fs a -> fs :/ FieldApp a
unExtensibleUnion
    {-# INLINE absurdUnion #-}

    inject0 :: forall (f :: * -> *) (fs :: [* -> *]).
f ~> ExtensibleUnion (f : fs)
inject0 = forall (fs :: [* -> *]) a.
(fs :/ FieldApp a) -> ExtensibleUnion fs a
ExtensibleUnion forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (xs :: [k]) (x :: k) (h :: k -> *).
Membership xs x -> h x -> xs :/ h
EmbedAt forall {k} (x :: k) (xs :: [k]). Membership (x : xs) x
leadership forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (f :: * -> *). f a -> FieldApp a f
FieldApp
    {-# INLINE inject0 #-}

    weaken :: forall (fs :: [* -> *]) a (f :: * -> *).
ExtensibleUnion fs a -> ExtensibleUnion (f : fs) a
weaken (ExtensibleUnion (EmbedAt Membership fs x
w FieldApp a x
e)) =
        forall (fs :: [* -> *]) a.
(fs :/ FieldApp a) -> ExtensibleUnion fs a
ExtensibleUnion forall a b. (a -> b) -> a -> b
$ forall {k} (xs :: [k]) (x :: k) (h :: k -> *).
Membership xs x -> h x -> xs :/ h
EmbedAt (forall {k} (xs :: [k]) (y :: k) (x :: k).
Membership xs y -> Membership (x : xs) y
nextMembership Membership fs x
w) FieldApp a x
e
    {-# INLINE weaken #-}

    f a -> r
f |+|: :: forall (f :: * -> *) a r (fs :: [* -> *]).
(f a -> r)
-> (ExtensibleUnion fs a -> r) -> ExtensibleUnion (f : fs) a -> r
|+|: ExtensibleUnion fs a -> r
g = (f a -> r
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a (f :: * -> *). FieldApp a f -> f a
unFieldApp forall {k} (h :: k -> *) (x :: k) r (xs :: [k]).
(h x -> r) -> ((xs :/ h) -> r) -> ((x : xs) :/ h) -> r
<:| ExtensibleUnion fs a -> r
g forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (fs :: [* -> *]) a.
(fs :/ FieldApp a) -> ExtensibleUnion fs a
ExtensibleUnion) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (fs :: [* -> *]) a. ExtensibleUnion fs a -> fs :/ FieldApp a
unExtensibleUnion
    {-# INLINE (|+|:) #-}

findFirstMembership :: forall xs x. KnownNat (TypeIndex xs x) => Membership xs x
findFirstMembership :: forall {k} (xs :: [k]) (x :: k).
KnownNat (TypeIndex xs x) =>
Membership xs x
findFirstMembership = Int -> Membership xs x
unsafeMkMembership forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal @(TypeIndex xs x) forall {k} (t :: k). Proxy t
Proxy
  where
    -- This hack may break if the membership package version gets updated.
    unsafeMkMembership :: Int -> Membership xs x
    unsafeMkMembership :: Int -> Membership xs x
unsafeMkMembership = forall a b. a -> b
unsafeCoerce

type family TypeIndex (xs :: [k]) (x :: k) :: Nat where
    TypeIndex (x ': xs) x = 0
    TypeIndex (y ': xs) x = 1 + TypeIndex xs x