Modular arithmetic with propagating type-checked modulus, Haskell

Modular arithmetic with propagating type-checked modulus

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ImpredicativeTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}

import Data.Word(Word8)
import Data.Bits
import System.IO.Unsafe
import Foreign.Storable
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array
import Foreign.Ptr
import Foreign.Marshal.Utils
import Data.ByteString(ByteString)
import Control.Applicative(Const)

data Zero;
data One;

data Unit a = Unit;

-- Reflecting a single bit
class ReflectBit s where
    reflectBit :: Unit s -> Word8

instance ReflectBit Zero where
    reflectBit _ = 0
instance ReflectBit One where
    reflectBit _ = 1
        
reifyBit :: Word8 -> (forall s. ReflectBit s => Unit s -> r) -> r
reifyBit b k = case b of
    0 -> k (Unit :: Unit Zero)
    otherwise -> k (Unit :: Unit One)

-- Reflecting a whole word at once
data WordT b0 b1 b2 b3 b4 b5 b6 b7;

class ReflectWord s where
    reflectWord :: Unit s -> Word8

instance (ReflectBit b0, ReflectBit b1, ReflectBit b2, ReflectBit b3,
    ReflectBit b4, ReflectBit b5, ReflectBit b6, ReflectBit b7) =>
    ReflectWord (WordT b0 b1 b2 b3 b4 b5 b6 b7) where
    reflectWord (_ :: Unit (WordT b0 b1 b2 b3 b4 b5 b6 b7))
        = (reflectBit (Unit :: Unit b0))
        + (reflectBit (Unit :: Unit b1)) * 2
        + (reflectBit (Unit :: Unit b2)) * 4
        + (reflectBit (Unit :: Unit b3)) * 8
        + (reflectBit (Unit :: Unit b4)) * 16
        + (reflectBit (Unit :: Unit b5)) * 32
        + (reflectBit (Unit :: Unit b6)) * 64
        + (reflectBit (Unit :: Unit b7)) * 128

reifyWord :: Word8 -> (forall s. ReflectWord s => Unit s -> r) -> r
reifyWord w k = reifyBit (w .&. 1) reifyWord0 where
    reifyWord0 l0 = reifyBit (w .&. 2) (reifyWord1 l0) where
    reifyWord1 l0 l1 = reifyBit (w .&. 4) (reifyWord2 l0 l1) where
    reifyWord2 l0 l1 l2 = reifyBit (w .&. 8) (reifyWord3 l0 l1 l2) where
    reifyWord3 l0 l1 l2 l3 = reifyBit (w .&. 16) (reifyWord4 l0 l1 l2 l3) where
    reifyWord4 l0 l1 l2 l3 l4 = reifyBit (w .&. 32) (reifyWord5 l0 l1 l2 l3 l4) where
    reifyWord5 l0 l1 l2 l3 l4 l5 = reifyBit (w .&. 64) (reifyWord6 l0 l1 l2 l3 l4 l5) where
    reifyWord6 l0 l1 l2 l3 l4 l5 l6 = reifyBit (w .&. 128) (reifyWord7 l0 l1 l2 l3 l4 l5 l6) where
    reifyWord7 (_ :: Unit b0) (_ :: Unit b1) (_ :: Unit b2) (_ :: Unit b3) (_ :: Unit b4) (_ :: Unit b5) (_ :: Unit b6) (_ :: Unit b7)= k (Unit :: Unit (WordT b0 b1 b2 b3 b4 b5 b6 b7))

-- Reflecting any [Word8]
data Nil;
data Cons s ss;

class ReflectBytes s where
    reflectBytes :: Unit s -> [Word8]

instance ReflectBytes Nil where
    reflectBytes _ = []
instance (ReflectWord b, ReflectBytes bs) => ReflectBytes (Cons b bs) where
    reflectBytes (_ :: Unit (Cons b bs)) = reflectWord (Unit :: Unit b) : reflectBytes (Unit :: Unit bs)

reifyBytes :: [Word8] -> (forall s. ReflectBytes s => Unit s -> r) -> r
reifyBytes [] k = k (Unit :: Unit Nil)
reifyBytes (b:bs) k = reifyWord b reifyBytes' where
    reifyBytes' w0 = reifyBytes bs (reifyBytes'' w0) where
    reifyBytes'' (_ :: Unit w) (_ :: Unit ws) = k (Unit :: Unit (Cons w ws))
    
-- Reflecting any Storable
data Store s a;

class ReflectStorable s where
    reflectStorable :: Storable a => Unit (s a) -> a

instance ReflectBytes b => ReflectStorable (Store b) where
    reflectStorable _ = unsafePerformIO
        $ alloca
        $ \s -> do pokeArray (castPtr s) $ reflectBytes (Unit :: Unit b)
                   peek s

reifyStorable :: Storable a => a -> (forall s. ReflectStorable s => Unit (s a) -> r) -> r
reifyStorable w k = reifyBytes bytes reifyStorable' where
    reifyStorable' (_ :: Unit s) = k (Unit :: Unit (Store s a))
    bytes = unsafePerformIO $ with w (peekArray (sizeOf w) . castPtr)

-- Using it for modular computations
newtype M s a = M a; -- our modular number type
unM :: M s a -> a
unM (M a) = a

class Modular s a | s -> a where
    modulus :: Unit s -> a

instance (Storable a, Integral a, ReflectStorable s)
    => Modular (s a) a where
    modulus = reflectStorable

normalize :: forall s a. (Modular s a, Integral a) => a -> M s a
normalize v = M $ mod v (modulus (Unit :: Unit s))

instance (Modular s a, Integral a) => Num (M s a) where
    (M a) + (M b) = normalize (a + b)
    (M a) - (M b) = normalize (a - b)
    (M a) * (M b) = normalize (a * b)
    abs (M a) = normalize $ abs a
    signum (M a) = normalize $ signum a
    fromInteger a = normalize $ fromInteger a

withModular :: (Storable a, Integral a) => a -> (forall s. Modular s a => M s a) -> a
withModular i k
    = reifyStorable i (\(_ :: Unit (s a)) -> unM (k :: M (s a) a))

-- example computation
exampleComputation :: (Integral a) => Modular s a => M s a
exampleComputation = 4 * 4 + 5 * 5 + 9 * 5 * 17827

main = putStrLn $ show $ (withModular (10001 :: Int) exampleComputation)


run \| edit \| history \| help	0