{-# 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 on the list is in WHNF (the spine may be lazy)
-- and the length of the list is the same as of the type-level shape.
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 function makes no attempt to stop you from breaking the data invariant for 'ListX'. If you do so, you must later ensure that the invariant is reinstated, for example using @'lazilyForce' 'id'@.
{-# INLINE lazily #-}
lazily :: ([a] -> [b]) -> ListX sh a -> ListX sh b
lazily f (ListX l) = ListX $ f l

-- | This function makes no attempt to stop you from breaking the data invariant for 'ListX'. If you do so, you must later ensure that the invariant is reinstated, for example using @'lazilyForce' 'id'@.
{-# 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 forces all elements of the @[b]@ list to restore the strictness part of the data invariant for 'ListX'. Note that ensuring the list has the right length is still the user's responsibility.
{-# 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