let it leave me like a long breath

i’m writing some haskell & i’ve put together maybe the first complicated, math-based type i’ve ever written?

it’s for enumerable combinatorics equations and the deal is you got

total :: Enumerable a -> Integer
pick :: Monoid a => Enumerable a -> Integer -> Either EnumerableError a

where you can say ‘give me anything in [0..total x)’ and you’ll get out that specific permutation of values. so it lets you exhaustively enumerate certain random generators, basically. and also transform the random generation into just "pick a random number in the range and then just calculate that value", since it's possible to generate any given value without looking at its neighbors.

(the monoid constraint is just b/c otherwise it would have to return [a] which is kind of annoying since that’s basically never what you’d actually want)

i was testing this with quickcheck (also my first time using quickcheck) to make sure it actually did what i thought it did, and for a while i had a quickcheck Arbitrary instance that just did not terminate, since these are nested data structures that can store an arbitrarily complex expression. even after limiting it a bunch i still can get it to churn out nice compact expressions that have a total count of, for example, 458575485777749657103020327287430331293288631270492527580669116448151889324257271211264327503524251371913200211982149838423029009393779536545995926888936184639883868071199119583128675780396729867476981269731795729169851133510432815404825306221507642742572989858640526486864585952735243739518209812574667794254182198110257592401920000000000000000000000000000000000000000000000000000000000000000000000000000000000

so it has kind of been illuminating about the nature of combinatorial explosions.

i’m working on adding weights and constraints and some other stuff but that’s somewhat difficult, since some constraints are basically super easy to codify and others are known NP-complete.

i spent a lot of time thinking about how it would or would not be an applicative functor, and then it turns out the actual Applicative instance is like ~250 characters

anyway code follows

module Data.Enumerable
	( EnumerableError(..)
	, Enumerable
	-- use functions
	, total
	, pick

	-- construction functions
	, range
	, seriesE
	, parallelE
	, oneE
	) where

import Data.Monoid
import Data.List (mapAccumL)

import Test.QuickCheck (Arbitrary(..), Gen, sized, oneof, resize, listOf1)

{-
maybe add some basic scanl/scanr functions for building Parallel/Series? or unfoldr. just "run this function on its output until it stops and collect all of those as options"
-}

data EnumerableError = UnderMin Integer | OverMax Integer
	deriving (Eq, Ord, Show, Read)

data Enumerable a
	= One a
	| Parallel Integer Integer Int [(Integer, Enumerable a)]
	| Series Integer Integer Int [Enumerable a]
	deriving (Eq, Ord, Show, Read)

instance Functor Enumerable where
	fmap f (One a) = One (f a)
	fmap f (Parallel a b c vs) = Parallel a b c $ fmap (\(i, v) -> (i, fmap f v)) vs
	fmap f (Series a b c vs) = Series a b c (fmap (fmap f) vs)

instance Foldable Enumerable where
	foldMap f (One a) = f a
	foldMap f (Parallel _ _ _ vs) = mconcat . fmap (foldMap f . snd) $ vs
	foldMap f (Series _ _ _ vs) = mconcat $ foldMap f <$> vs

instance Traversable Enumerable where
	traverse f (One a) = One <$> f a
	traverse f (Parallel a b c vs) = recount . Parallel a b c <$> (sequenceA $ (\(i, v) -> (,) i <$> traverse f v) <$> vs)
	traverse f (Series a b c vs) = recount . Series a b c <$> (sequenceA . fmap (traverse f) $ vs)

instance Applicative Enumerable where
	pure = One
	(One f) <*> v = f <$> v
	(Parallel a b c fs) <*> v = recount $ Parallel a b c $ (\(i, f) -> (i, f <*> v)) <$> fs
	-- i don't know if a Series can ever actually get 'longer' but i'd rather not risk it, since a Parallel certainly can
	(Series a b c fs) <*> v = recount $ Series a b c $ (\f -> f <*> v) <$> fs

range :: Enum a => (a, a) -> Enumerable a
range (lo, hi) = parallelE $ oneE <$> [lo..hi]

-- quickcheck instances
instance Arbitrary a => Arbitrary (Enumerable a) where
	arbitrary = sized sizedArbitrary
	shrink (One _) = []
	shrink (Parallel _ _ _ subs) = fmap snd subs <> ((shrink . snd) =<< subs)
	shrink (Series _ _ _ subs) = subs <> (shrink =<< subs)

sizedArbitrary :: Arbitrary a => Int -> Gen (Enumerable a)
sizedArbitrary 0 = oneE <$> arbitrary
sizedArbitrary 1 = oneE <$> arbitrary
sizedArbitrary n = oneof
	[ oneE <$> arbitrary
	, parallelE <$> resize limit (listOf1 (sizedArbitrary limit))
	, seriesE <$> resize limit (listOf1 (sizedArbitrary limit))
	]
	where
		limit = n `div` 2 -- round $ sqrt $ fromIntegral n



recount :: Enumerable a -> Enumerable a
recount = snd . prepareCount 1

countOptions :: [Enumerable a] -> (Integer, [Enumerable a])
countOptions = mapAccumL prepareCount 1

prepareCount :: Integer -> Enumerable a -> (Integer, Enumerable a)
prepareCount c (One x) = (c, One x)
-- a Parallel is different from a Series: a series evaluates every option in its list in order, whereas a parallel picks ONE option from its list
prepareCount c (Parallel _ _ len vs) =
	let (subopts, vs') = mapAccumL count 0 . fmap (countOptions . pure) $ fmap snd vs
	in (c * subopts, Parallel (c * subopts) c len vs')
	where
		count :: Integer -> (Integer, [Enumerable a]) -> (Integer, (Integer, Enumerable a))
		count cur (next, [sel]) = (cur + next, (cur, sel))
		count cur (next, _) = error "prepareCount incoherent"
prepareCount c (Series _ _ len xs) =
	let (subopts, xs') = countOptions xs
	in (c * subopts, Series (c * subopts) c len xs')

seriesE :: [Enumerable a] -> Enumerable a
seriesE vs = recount $ Series 1 1 (length vs) vs

parallelE :: [Enumerable a] -> Enumerable a
parallelE vs = recount $ Parallel 1 1 (length vs) ((\x -> (0, x))<$> vs)

oneE :: a -> Enumerable a
oneE v = One v

total :: Enumerable a -> Integer
total (One _) = 1
total (Parallel opts _ _ _) = opts
total (Series opts _ _ _) = opts

pick :: Monoid a => Enumerable a -> Integer -> Either EnumerableError a
pick (One v) 0 = Right v
pick (One v) _ = Left $ OverMax 1
pick many@(Parallel opts _ _ vs) x
	| x < 0 = Left $ UnderMin 0
	| x >= opts = Left $ OverMax opts
	| otherwise = Right $ mconcat $ select x many
pick (Series opts _ _ vs) x
	| x < 0 = Left $ UnderMin 0
	| x >= opts = Left $ OverMax opts
	| otherwise = Right $ mconcat $ pick_ x vs

pick_ :: Integer -> [Enumerable a] -> [a]
pick_ i = (select i =<<)

select :: Integer -> Enumerable a -> [a]
select i (One x) = [x]
select i (Parallel opts cycle _ xs) =
	let j = (i `mod` opts) `div` cycle
	in
		-- we want the last value that's smaller than the computed index; this unfortunately means using `last` right now
		case takeWhile (\(k, _) -> j >= k) xs of
			[] -> error "select ran out of values while picking, but didn't hit the max index count. ??"
			-- :(
			vs -> (\(k, val) -> select (j - k) val) $ last vs
select i (Series opts cycle _ xs) = pick_ ((i `mod` opts) `div` cycle) xs