A type for doing modulo arithmetic that forces all calculations to use the same mod at a type level, but allows the mod to be decided by a value

ModArith.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleInstances #-}

module ModArith
	( WithMod
	, IntMod
	, doMod
	) where

import Control.Applicative (liftA2)

newtype WithMod n x = WithMod { getMod :: Integer -> x }

instance Functor (WithMod n) where
	fmap f (WithMod g) = WithMod $ f <$> g

instance Applicative (WithMod n) where
	pure = WithMod . pure
	WithMod f <*> WithMod g = WithMod $ f <*> g

instance Monad (WithMod n) where
	WithMod f >>= g = WithMod $ \x -> (getMod $ g $ f x) x

doMod :: Integer -> (forall n. WithMod n x) -> x
doMod i (WithMod f) = f i

type IntMod n = WithMod n Integer

-- Force an IntMod into the range 0..n-1
normalise :: IntMod n -> IntMod n
normalise (WithMod f) = WithMod $ \x -> f x `mod` x

instance Num (WithMod n Integer) where
	(+) = (normalise .) . liftA2 (+)
	(-) = (normalise .) . liftA2 (-)
	(*) = (normalise .) . liftA2 (*)
	negate = normalise . fmap negate
	abs = id
	signum = id
	fromInteger = normalise . pure