Lists with type specified length

FixedList.hs

{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module FixedList
	( FixedList(unFix)
	, toFixed
	, nil
	, cons
	, (+:)
	, (+++)
	, head
	, last
	, tail
	, init
	, uncons
	, reverse
	, intersperse
	, intercalate
	, transpose
	, concat
	, scanl
	, scanl1
	, scanr
	, scanr1
	, iterate
	, unfoldr
	, take
	, drop
	, split
	, sort
	, sortOn
	, insert
	, sortBy
	, insertBy
	) where

import GHC.TypeNats
	(type (-), type (<=), type (+), type (*), Nat, natVal, KnownNat)
import qualified Data.List as L
import Data.Proxy (Proxy(Proxy))
import Data.Function (on)
import Data.Functor.Classes (Eq1, Ord1, Show1(liftShowsPrec), showsPrec1)

import Prelude hiding
	( head, tail, init, last, reverse, concat, scanl, scanl1, scanr, scanr1
	, iterate, take, drop
	)

-- Helper for using a Nat in take, drop, replicate, and split
toInt :: forall n. KnownNat n => Int
toInt = fromInteger $ toInteger $ natVal (Proxy :: Proxy n)

-- A list of length n containing elements of type a
newtype FixedList (n :: Nat) a = FixedList { unFix :: [a] }
	deriving (Eq, Ord, Eq1, Ord1)

instance Show a => Show (FixedList n a) where
	showsPrec = showsPrec1

instance Show1 (FixedList n) where
	liftShowsPrec f _ p (FixedList l) = showParen (p > 5) $
		foldr
			(\a b -> f 6 a . showString " +: " . b)
			(showString "nil")
			l

instance Functor (FixedList n) where
	fmap f = FixedList . fmap f . unFix

-- Applicative instance based on zip lists to preserve length
instance KnownNat n => Applicative (FixedList n) where
	pure = FixedList . replicate (toInt @n)
	FixedList a <*> FixedList b = FixedList $ zipWith ($) a b

instance Foldable (FixedList n) where
	foldMap f = foldMap f . unFix

instance Traversable (FixedList n) where
	sequenceA = fmap FixedList . sequenceA . unFix

-- Get the first n elements of a list, fails if list is too short
toFixed :: forall a n. KnownNat n => [a] -> Maybe (FixedList n a)
toFixed l
	| length result == toInt @n = Just $ FixedList result
	| otherwise = Nothing
	where
		result = L.take (toInt @n) l

-- The list of length 0
nil :: FixedList 0 a
nil = FixedList []

-- Add an element to the front of a list
infixr 5 `cons`
cons :: a -> FixedList n a -> FixedList (n+1) a
cons x (FixedList xs) = FixedList $ x:xs

-- Operator version of cons
infixr 5 +:
(+:) = cons

-- Concatenate two lists
infixr 5 +++
(+++) :: FixedList n a -> FixedList m a -> FixedList (n+m) a
FixedList a +++ FixedList b = FixedList $ a ++ b

-- Get the first element of a non-empty list
head :: 1 <= n => FixedList n a -> a
head = L.head . unFix

-- Get the last element of a non-empty list
last :: 1 <= n => FixedList n a -> a
last = L.last . unFix

-- Get the last elements of a non-empty list
tail :: 1 <= n => FixedList n a -> FixedList (n-1) a
tail = FixedList . L.tail . unFix

-- Get the first elements of a non-empty list
init :: 1 <= n => FixedList n a -> FixedList (n-1) a
init = FixedList . L.init . unFix

-- Get the first element and the rest of a non-empty list
uncons :: 1 <= n => FixedList n a -> (a, FixedList (n-1) a)
uncons (FixedList (x:xs)) = (x, FixedList xs)
uncons _ = error "Invalid FixedList"

-- Reverse a list
reverse :: FixedList n a -> FixedList n a
reverse = FixedList . L.reverse . unFix

-- Make a list alternating between the given list and the given element
intersperse :: 1 <= n => a -> FixedList n a -> FixedList (2*n - 1) a
intersperse c (FixedList l) = FixedList $ L.intersperse c l

-- Concatenate all the given lists inserting the given list between them
intercalate :: 1 <= y =>
	FixedList x a -> FixedList y (FixedList z a) -> FixedList (y*z + x*(y-1)) a
intercalate (FixedList a) (FixedList b) =
	FixedList $ L.intercalate a $ unFix <$> b

-- Swap rows and columns
transpose :: FixedList n (FixedList m a) -> FixedList m (FixedList n a)
transpose = FixedList . fmap FixedList . L.transpose . fmap unFix . unFix

-- Concatenate all the lists
concat :: FixedList n (FixedList m a) -> FixedList (n*m) a
concat = FixedList . concatMap unFix . unFix

-- Make a list starting with the given element and applying the given function
-- to the last element produced and the first unused element of the given list
-- to make the next element of the produced list
scanl :: (b -> a -> b) -> b -> FixedList n a -> FixedList (n+1) b
scanl f s = FixedList . L.scanl f s . unFix

-- Uses the first element of the list as the starting element in scanl
scanl1 :: (a -> a -> a) -> FixedList n a -> FixedList n a
scanl1 f = FixedList . L.scanl1 f . unFix

-- Does scanl in reverse order
scanr :: (a -> b -> b) -> b -> FixedList n a -> FixedList (n+1) b
scanr f s = FixedList . L.scanr f s . unFix

-- Does scanl1 in reverse order
scanr1 :: (a -> a -> a) -> FixedList n a -> FixedList n a
scanr1 f = FixedList . L.scanr1 f . unFix

-- Makes a list starting with the given element and each successive element
-- is the result of applying the function to the element before
iterate :: forall a n. KnownNat n => (a -> a) -> a -> FixedList n a
iterate f s = FixedList $ L.take (toInt @n) $ L.iterate f s

-- Makes a list by applying the function to the accumulator to get the
-- next list element and the new accumulator
unfoldr :: forall a b n. KnownNat n => (b -> (a,b)) -> b -> FixedList n a
unfoldr f s = FixedList $ L.take (toInt @n) $ L.unfoldr (Just . f) s

-- Keeps the start of the list
take :: forall a m n. (KnownNat n, n <= m) => FixedList m a -> FixedList n a
take = FixedList . L.take (toInt @n) . unFix

-- Keeps the end of the list
drop :: forall a m n. (KnownNat (m-n), n <= m) => FixedList m a -> FixedList n a
drop = FixedList . L.drop (toInt @(m-n)) . unFix

-- Splits a list in two
split :: forall a m n. KnownNat m =>
	FixedList (m+n) a -> (FixedList m a, FixedList n a)
split (FixedList l) = (FixedList a, FixedList b) where
	(a,b) = splitAt (toInt @m) l

-- Sorts a list
sort :: Ord a => FixedList n a -> FixedList n a
sort = FixedList . L.sort . unFix

-- Sorts a list using the function to decide the order of elements
sortOn :: Ord b => (a -> b) -> FixedList n a -> FixedList n a
sortOn f = FixedList . L.sortOn f . unFix

-- Inserts a value before the first element bigger than it
insert :: Ord a => a -> FixedList n a -> FixedList (n+1) a
insert e = FixedList . L.insert e . unFix

-- Sorts a list using the function to decide the order of elements
sortBy :: (a -> a -> Ordering) -> FixedList n a -> FixedList n a
sortBy f = FixedList . L.sortBy f . unFix

-- Inserts a list before the first element bigger than it according to the
-- given function
insertBy :: (a -> a -> Ordering) -> a -> FixedList n a -> FixedList n a
insertBy f e = FixedList . L.insertBy f e . unFix