aboutsummaryrefslogtreecommitdiff
path: root/src/Data/Array/Nested/Mixed/ListX.hs
blob: e4e224bef06c6a12b070e7bf7ac2e011ed027884 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
{-# 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