{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.Array.Nested.Mixed.ListX (ListX, pattern ZX, pattern (::%), listxShow, lazily, lazilyConcat, lazilyForce, Rank, coerceEqualRankListX) where

import Control.DeepSeq (NFData(..))
import Data.Foldable qualified as Foldable
import Data.Kind (Type)
import Data.Type.Equality
import GHC.IsList (IsList)
import GHC.IsList qualified as IsList
import GHC.TypeLits
#if !MIN_VERSION_GLASGOW_HASKELL(9,8,0,0)
import GHC.TypeLits.Orphans ()
#endif

import Data.Array.Nested.Types


-- | The length of a type-level list. If the argument is a shape, then the
-- result is the rank of that shape.
type family Rank sh where
  Rank '[] = 0
  Rank (_ : sh) = Rank sh + 1


-- * Mixed lists implementation

-- | Data invariant: each element is in WHNF (the spine may be be not forced).
type role ListX nominal representational
type ListX :: [Maybe Nat] -> Type -> Type
newtype ListX sh i = ListX [i]
  deriving (Eq, Ord, NFData, Foldable)

{-# INLINE ZX #-}
pattern ZX :: forall sh i. () => sh ~ '[] => ListX sh i
pattern ZX <- (listxNull -> Just Refl)
  where ZX = ListX []

{-# INLINE listxNull #-}
listxNull :: ListX sh i -> Maybe (sh :~: '[])
listxNull (ListX []) = Just unsafeCoerceRefl
listxNull (ListX (_ : _)) = Nothing

{-# INLINE (::%) #-}
pattern (::%)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => i -> ListX sh i -> ListX sh1 i
pattern i ::% l <- (listxUncons -> Just (UnconsListXRes i l))
  where !i ::% ListX !l = ListX (i : l)
infixr 3 ::%

data UnconsListXRes i sh1 =
  forall n sh. (n : sh ~ sh1) => UnconsListXRes i (ListX sh i)
{-# INLINE listxUncons #-}
listxUncons :: forall sh1 i. ListX sh1 i -> Maybe (UnconsListXRes i sh1)
listxUncons (ListX (i : l)) = gcastWith (unsafeCoerceRefl :: Head sh1 ': Tail sh1 :~: sh1) $
                              Just (UnconsListXRes i (ListX @(Tail sh1) l))
listxUncons (ListX []) = Nothing

{-# COMPLETE ZX, (::%) #-}

-- | This operation may not preserve the data invariant for ListX,
-- regardless whether the invariant holds for the argument list.
-- It's the user's obligation to honor the data invariant,
-- e.g., by forcing all new elements of the resulting list at the end.
{-# INLINE lazily #-}
lazily :: ([a] -> [b]) -> ListX sh a -> ListX sh b
lazily f (ListX l) = ListX $ f l

-- | This operation may not preserve the data invariant for ListX,
-- regardless whether the invariant holds for the argument lists.
-- It's the user's obligation to honor the data invariant,
-- e.g., by forcing all new elements of the resulting list at the end.
{-# INLINE lazilyConcat #-}
lazilyConcat :: ([a] -> [b] -> [c]) -> ListX sh a -> ListX sh' b -> ListX (sh ++ sh') c
lazilyConcat f (ListX l) (ListX k) = ListX $ f l k

-- | This operation honours the data invariant for @ListX@ by forcing
-- all elements of the resulting list at the end.
{-# INLINE lazilyForce #-}
lazilyForce :: ([a] -> [b]) -> ListX sh a -> ListX sh b
lazilyForce f (ListX l) = let res = f l
                          in foldr seq () res `seq` ListX res

#ifdef OXAR_DEFAULT_SHOW_INSTANCES
deriving instance Show i => Show (ListX sh i)
#else
instance Show i => Show (ListX sh i) where
  showsPrec _ = listxShow shows
#endif

{-# INLINE listxShow #-}
listxShow :: forall sh i. (i -> ShowS) -> ListX sh i -> ShowS
listxShow f l = showString "[" . go "" l . showString "]"
  where
    go :: String -> ListX sh' i -> ShowS
    go _ ZX = id
    go prefix (x ::% xs) = showString prefix . f x . go "," xs

-- This can't be derived, becauses the list needs to be fully evaluated,
-- per data invariant. This version is faster than versions defined using
-- (::%) or lazilyForce.
instance Functor (ListX l) where
  {-# INLINE fmap #-}
  fmap f (ListX l) =
    let fmap' [] = []
        fmap' (x : xs) = let y = f x
                             rest = fmap' xs
                         in y `seq` rest `seq` (y : rest)
    in ListX $ fmap' l

-- | Very untyped: not even length is checked (at runtime).
instance IsList (ListX sh i) where
  type Item (ListX sh i) = i
  {-# INLINE fromList #-}
  fromList l = foldr seq () l `seq` ListX l
  {-# INLINE toList #-}
  toList = Foldable.toList

coerceEqualRankListX :: Rank sh ~ Rank sh' => ListX sh i -> ListX sh' i
coerceEqualRankListX (ListX l) = ListX l