icfp13/Encoding.hs

149 wiersze
3.7 KiB
Haskell

{-# OPTIONS -fglasgow-exts #-}
module Encoding (Trit(..), Encode(..), Stream(..), encode, decode, decodeStrict, toTrit, fromTrit) where
import Control.Monad
data Trit = T0 | T1 | T2 deriving (Enum, Eq)
instance Show Trit where
show = show . fromEnum
instance Read Trit where
readsPrec _ "0" = [(T0, "")]
readsPrec _ "1" = [(T1, "")]
readsPrec _ "2" = [(T2, "")]
instance Num Trit where
(+) = liftToTrit (+)
(*) = liftToTrit (*)
(-) = liftToTrit (-)
abs = id
signum = const 0
fromInteger 0 = T0
fromInteger 1 = T1
fromInteger 2 = T2
liftToTrit :: (Int -> Int -> Int) -> Trit -> Trit -> Trit
liftToTrit f x y = toEnum $ (fromEnum x `f` fromEnum y) `mod` 3
toTrit :: Int -> Trit
toTrit = toEnum
fromTrit :: Integral i => Trit -> i
fromTrit = fromIntegral . fromEnum
class Stream a where
streamread :: a -> [Trit]
streamwrite :: [Trit] -> a
instance Stream [Trit] where
streamread = id
streamwrite = id
instance Stream [Int] where
streamread = map toEnum
streamwrite = map fromEnum
instance Stream String where
streamread = map (toEnum . read . (:[]))
streamwrite = concat . map (show . fromEnum)
data Decoder i o = Decoder { doDecode :: [i] -> ([i], o) }
type Encoder o i = i -> [o]
type DR a = Decoder Trit a
instance Monad (Decoder i) where
f >>= g = Decoder $ \x -> let (x0, r0) = doDecode f x in doDecode (g r0) x0
f >> g = Decoder $ \x -> let (x0, r0) = doDecode f x in doDecode g x0
return a = Decoder $ \x -> (x, a)
fail = error
class Encode a where
sdecode :: DR a
sencode :: Encoder Trit a
encode :: (Stream s, Encode a) => a -> s
encode = streamwrite . sencode
decode :: (Stream s, Encode a) => s -> a
decode = snd . (doDecode sdecode) . streamread
decodeStrict :: (Stream s, Encode a) => s -> a
decodeStrict x = case (doDecode sdecode $ streamread x) of
([], a) -> a
(_, _) -> error "Expected end of input"
getTrit :: Decoder Trit Trit
getTrit = Decoder $ \xs -> case xs of
[] -> error "Unexpected end of stream"
(i:xs) -> (xs, i)
instance Encode Bool where
sdecode = do
i <- sdecode :: DR Int
case i of
0 -> return False
1 -> return True
_ -> fail "Expected 0 or 1"
sencode True = [1]
sencode False = [0]
instance Encode Trit where
sdecode = getTrit
sencode t = [t]
instance Encode a => Encode [a] where
sdecode = do
i <- getTrit
case i of
T0 -> return []
T1 -> sdecode >>= (\x -> return [x])
T2 -> do
i <- getTrit
case i of
T2 -> do
n <- sdecode :: DR Int
replicateM (n+2) sdecode
_ -> fail "Unexpected end of stream"
sencode [] = [T0]
sencode [x] = T1:sencode x
sencode xs = T2:T2:(sencode $ length xs - 2) ++ concat (map sencode xs)
instance Encode Int where
sdecode = do
n <- sdecode :: DR Integer
return $ fromIntegral n
sencode i = sencode (fromIntegral i :: Integer)
instance Encode Integer where
sdecode = do
digits <- sdecode :: DR [Trit]
return $ decodeTern digits + (3^(length digits - 1) `div` 2)
sencode n
| n < 0 = error "Can't encode negative numbers"
| n == 0 = [0]
| n > 0 = let len = log3 n in sencode (encodeTern len (n - base len)) where
base len = (3^len - 1) `div` 2
log3 n = (head $ filter (\len -> n < base len) [0..]) - 1
instance (Encode a, Encode z) => Encode (a, z) where
sdecode = do
a <- sdecode
z <- sdecode
return (a, z)
sencode (a, z) = sencode a ++ sencode z
instance (Encode (a, b), Encode z) => Encode (a, b, z) where
sdecode = do
(a, b) <- sdecode
z <- sdecode
return (a, b, z)
sencode (a, b, z) = sencode (a, b) ++ sencode z
decodeTern :: Integral i => [Trit] -> i
decodeTern digits = dec 0 digits where
dec s [] = s
dec s (i:xs) = dec (3*s + (fromTrit i)) xs
encodeTern :: Integral i => Int -> i -> [Trit]
encodeTern 0 0 = []
encodeTern k x = (encodeTern (k-1) (x `div` 3)) ++ [toTrit $ fromIntegral (x `mod` 3)]