Bin.idr

||| Binary parsing, based on the stuff in Power of Pi
module Bin


import System

import Data.So
import Data.Vect
import Bytes

%default total

data BaseData : Type where
  ||| A byte
  Byte : BaseData
  ||| n repetitions of some base data
  NTimes : Nat -> BaseData -> BaseData
  ||| An ASCII Char
  Ch : BaseData

||| Interpret the base data as a type
baseType : BaseData -> Type
baseType Byte = Bits8
baseType (NTimes n b) = Vect n (baseType b)
baseType Ch = Char

||| A tagged bit of code, to make things easier to read and debug
data Section : String -> Type -> Type where
  MkSection : (lbl : String) -> (x : a) -> Section lbl a

mutual
  ||| Format specifiers, which combine base data
  data Fmt : Type where
    Part : String -> Fmt -> Fmt
    Bad : String -> Fmt
    End : Fmt
    Base : BaseData -> Fmt
    OneOf : Fmt -> Fmt -> Fmt
    Many : Fmt -> Fmt
    Skip : (fmt : Fmt) -> (rep : resultOf fmt) -> Fmt -> Fmt
    (>>=) : (fmt : Fmt) -> (resultOf fmt -> Fmt) -> Fmt

  ||| Compute the type that corresponds to base data
  resultOf : Fmt -> Type
  resultOf (Part lbl fmt) = Section lbl (resultOf fmt)
  resultOf (Bad _) = Void
  resultOf End = ()
  resultOf (Base b) = baseType b
  resultOf (OneOf x y) = Either (resultOf x) (resultOf y)
  resultOf (Many fmt) = List (resultOf fmt)
  resultOf (Skip _ _ y) = resultOf y
  resultOf (fmt >>= f) = (x : resultOf fmt ** resultOf (f x))

||| Match data that satisfies a predicate
satisfy : [static] String -> [static] (fmt : Fmt) -> (resultOf fmt -> Bool) -> Fmt
satisfy what fmt p =
  do x <- fmt
     if p x
       then End
       else Bad what

||| Match one particular byte
byte : Bits8 -> Fmt -> Fmt
byte b fmt = Skip (satisfy "byte" (Base Byte) (b ==))
                  (b ** rewrite primEqReflexive b in ()) fmt
  where postulate primEqReflexive : (x : Bits8) -> prim__eqB8 x x = 1

||| Match one byte that's an ASCII space or newline
whitespace : Fmt
whitespace = satisfy "whitespace" (Base Byte) (\x => x == 0x0a || x == 0x20)

||| Match at least one whitespace byte
whitespaces : Fmt
whitespaces = Skip whitespace (0x20 ** ()) (Many whitespace)

||| Match at least one ASCII digit
digit : Fmt
digit = satisfy "digit" (Base Ch) isDigit


data Result : Fmt -> Type where
  Failure : (msg : String) -> Result fmt
  Success : resultOf fmt -> Int -> Result fmt

||| Read the base datatypes from a byte array
|||
||| @ b the format to read
||| @ i the pointer to start at
||| @ j the size of the array
covering
readBase : [static] (b : BaseData) -> (i : Int) -> ByteArray j -> IO (Maybe (baseType b, Int))
readBase Byte i arr with (choose $ inBounds i j)
  readBase Byte i arr | Right x = return Nothing
  readBase Byte i arr | Left ok = do byte <- read i arr ok
                                     return (Just (byte, i+1))
readBase (NTimes n b) i arr = ntimes n i
  where partial
        ntimes : (n : Nat) -> Int -> IO (Maybe (Vect n (baseType b), Int))
        ntimes Z i = return (Just ([], i))
        ntimes (S k) i = do Just (x, i') <- readBase b i arr
                              | Nothing => return Nothing
                            Just (xs, i'') <- ntimes k i'
                              | Nothing => return Nothing
                            return $ Just (x::xs, i'')
readBase Ch i arr with (choose $ inBounds i j)
  readBase Ch i arr | Right x = return Nothing
  readBase Ch i arr | Left ok = do byte <- read i arr ok
                                   return (Just (chr (prim__zextB8_Int byte), i+1))


||| Read from a byte array according to some format
|||
||| @ fmt the format to read
||| @ i the pointer to start at
||| @ j the size of the array
covering
readFmt : [static] (fmt : Fmt) -> (i : Int) -> ByteArray j -> IO (Result fmt)
readFmt (Part lbl fmt)     i arr =
  putStrLn lbl >>= \_ =>
  case !(readFmt fmt i arr) of
    Failure msg => return (Failure $ "In " ++ lbl ++ ": " ++ msg)
    Success x i' => return (Success (MkSection lbl x) i')
readFmt (Bad msg)          i arr = return (Failure $ "Couldn't get: " ++ msg)
readFmt End                i arr = return (Success () i)
readFmt (Base b)           i arr =
  case !(readBase b i arr) of
    Nothing        => return (Failure "couldn't read base")
    Just (elt, i') => return (Success elt i')
readFmt (OneOf fmt fmt')   i arr =
  case !(readFmt fmt i arr) of
    Failure msg => case !(readFmt fmt' i arr) of
                      Failure msg'   => return (Failure $ msg ++ msg')
                      Success res i' => return $ Success (Right res) i'
    Success res i' => return $ Success (Left res) i'
readFmt (Many fmt) i arr =
  case !(readFmt fmt i arr) of
    Failure _ => return $ Success [] i
    Success x i' =>
      case !(readFmt (Many fmt) i' arr) of
        Failure _ => return $ Success [x] i'
        Success xs i'' => return $ Success (x::xs) i''
readFmt (Skip fmt _ fmt') i arr =
  case !(readFmt fmt i arr) of
    Failure msg  => return $ Failure msg
    Success _ i' =>
      case !(readFmt fmt' i' arr) of
        Failure msg => return $ Failure msg
        Success res i'' => return $ Success res i''
readFmt (fmt >>= f)        i arr =
  case !(readFmt fmt i arr) of
    Failure msg => return $ Failure msg
    Success x i' => case !(readFmt (f x) i' arr) of
                      Failure msg => return $ Failure msg
                      Success y i'' => return $ Success (x ** y) i''

||| Read according to a format, starting with the beginning of the array
covering
readArray : [static] (fmt : Fmt) -> ByteArray j -> IO (Result fmt)
readArray fmt arr = readFmt fmt 0 arr


||| PPM raw images (no support for comments in the header yet)
ppm : Fmt
ppm = do Part "Magic" $ do
           byte 0x50 $ byte 0x36 $ whitespaces
         x <- Many digit
         y <- Skip whitespaces [(0x20 ** ())] $ Many digit
         let width  = dimension $ map getWitness x
         let height = dimension $ map getWitness y
         maxColor <- Skip whitespaces [(0x20 ** ())] $ Many digit
         let colspace = dimension $ map getWitness maxColor
         Skip whitespace (0x20 ** ()) $
           (Part "Image data:" $ Base $ NTimes width $ NTimes height $ NTimes 3 Byte)
         End
  where dimension : List Char -> Nat
        dimension xs = fromInteger $ cast $ pack xs

||| Flip the contents of PPM image data
mirror : resultOf ppm -> resultOf ppm
mirror (magic ** (width ** (height ** (maxColor ** (MkSection "Image data:" pixels ** ()))))) =
  (magic ** (width ** (height ** (maxColor ** (MkSection "Image data:" (reverse pixels) ** ())))))

||| Write base data to a byte array
covering -- TODO: make helper structurally recursive on format
writeBase : (b : BaseData) -> baseType b -> ByteArray k -> Int -> IO Int
writeBase {k} Byte x arr i with (choose $ inBounds i k)
  writeBase {k} Byte x arr i | Right y = do putStrLn "Overflow1" ; exit 1; return (-100)
  writeBase {k} Byte x arr i | Left ok = do write i x arr ok; return (i + 1)
writeBase (NTimes Z y) [] arr i = return i
writeBase (NTimes (S j) y) (x :: xs) arr i = do i' <- writeBase y x arr i
                                                writeBase (NTimes j y) xs arr i'
writeBase {k} Ch x arr i with (choose $ inBounds i k)
  writeBase {k = k} Ch x arr i | Right y = do putStrLn "Overflow2" ; exit 1; return (-100)
  writeBase {k = k} Ch x arr i | Left ok =
    do let byte = prim__truncInt_B8 $ prim__charToInt x
       write i byte arr ok
       return (i + 1)

||| Write formatted data to a byte array, starting at position 0
|||
||| @ fmt the format to follow
covering
writeArr : resultOf fmt -> ByteArray k -> IO ()
writeArr {k} res arr = do writeArr' _ res 0; return ()
  where partial
        writeArr' : (fmt : Fmt) -> resultOf fmt -> Int -> IO Int
        writeArr' (Part _ fmt) (MkSection _ res) i = writeArr' fmt res i
        writeArr' (Bad x) res i = absurd res
        writeArr' End () i = return i
        writeArr' (Base x) res i = writeBase x res arr i
        writeArr' (OneOf fmt _) (Left x)  i = writeArr' fmt x i
        writeArr' (OneOf _ fmt) (Right x) i = writeArr' fmt x i
        writeArr' (Many x) [] i = return i
        writeArr' (Many x) (y :: xs) i = writeArr' x y i >>= writeArr' (Many x) xs
        writeArr' (Skip x rep y) res i = writeArr' x rep i >>= writeArr' y res
        writeArr' (fmt >>= f) (x ** y) i = do i' <- writeArr' fmt x i
                                              writeArr' (f x) y i'

namespace Main
  partial
  main : IO ()
  main = do putStrLn "Running"
            Right (k ** image) <- slurp "me3.ppm"
              | Left err => do putStrLn err; exit 1
            putStrLn "Array slurped."
            res <- readArray ppm image
            putStrLn "Processed."
            case res of
              Failure msg => putStrLn $ "Failure: " ++ msg
              Success res i => do putStrLn $ "Success at " ++ show i ++ "!"
                                  arr' <- allocBytes (10 + k) (const 0)
                                  writeArr (mirror res) arr'
                                  dump "me3Flip.ppm" arr'