1 bit images in Haskell

Bitlang.hs

{-# LANGUAGE NumericUnderscores #-}

module Bitlang
    ( compile
    , evaluate
    , generate
    , ImageMeta(..)
    , ExprMeta(..)
    , Depth(..)
    ) where

import GHC.Conc                (getNumProcessors)
import System.IO               (stdout)
import Data.List               (words)
import Data.Array              (array, (!))
import Data.Maybe              (listToMaybe)
import Data.Monoid             ((<>))
import Control.Monad           (forM_, forever)
import Data.ByteString.Builder (word64BE, word32BE, hPutBuilder)

import Data.Bits
import Data.Word
import Control.Concurrent

newtype Expression = Expression [Item]
                     deriving (Show, Read)

data Depth = One
           | Sixteen
           deriving (Show, Read)

data Item = Add
          | Sub
          | Mul
          | Div
          | Mod
          | ShRight
          | ShLeft
          | Eq
          | Gt
          | Lt
          | Ge
          | Le
          | Or
          | And
          | Not
          | BitOr
          | BitAnd
          | BitXor
          | BitNot
          | Abs
          | VarX
          | VarY
          | VarW
          | VarH
          | VarT
          | Dup
          | Num NumT
          deriving (Show, Read)

type NumT = Int

data ImageMeta = ImageMeta { getWidth  :: NumT
                           , getHeight :: NumT
                           }
                 deriving (Show, Read)

data ExprMeta = ExprMeta { getW :: NumT
                         , getH :: NumT
                         , getX :: NumT
                         , getY :: NumT
                         , getT :: NumT
                         }
                deriving (Show, Read)

generate :: Depth -> Expression -> ImageMeta -> IO (Either String ())
generate depth expr imageMeta = do
    -- create thread pool to eval expressions in parallel
    nProcs <- getNumProcessors
    threadIn <- threadPool nProcs $
        \(mvar, meta, expr) -> do
            let x = getX meta
            let y = getY meta
            putMVar mvar $ evaluate meta expr

    let width = getWidth imageMeta
    let height = getHeight imageMeta

    arrayValues <- sequence $ take (fromIntegral $ width * height) $ repeat newEmptyMVar

    let indices = [ (x, y) | x <- [0..width-1], y <- [0..height-1] ]
    let results = array ((0,0), (width-1,height-1)) (indices `zip` arrayValues)

    -- spawn workers
    forkIO $ forM_ indices $
        \(x, y) -> do
            let mvar = results ! (x, y)
            let meta = ExprMeta width height x y 0
            writeChan threadIn (mvar, meta, expr)

    -- write farbfeld header
    putStr "farbfeld"
    hPutBuilder stdout $ word32BE $ fromIntegral width
    hPutBuilder stdout $ word32BE $ fromIntegral height

    -- write results
    let getPixel = case depth of
                     One     -> transformOneBit
                     Sixteen -> transformSixteenBits

    forM_ indices $
        \(x, y) -> do
            let mvar = results ! (x, y)
            result <- takeMVar mvar

            case result of
                Right value -> do
                    let encoded = word64BE $ getPixel value
                    hPutBuilder stdout encoded
                    return $ Right ()
                Left err -> return $ Left err

    return $ Right ()

transformOneBit :: NumT -> Word64
transformOneBit x = (y `shiftL` 48) .|. (y `shiftL` 32) .|. (y `shiftL` 16) .|. 0xffff
    where y = 0xffff * (fromIntegral x .&. 1)

transformSixteenBits :: NumT -> Word64
transformSixteenBits x = (y `shiftL` 48) .|. (y `shiftL` 32) .|. (y `shiftL` 16) .|. 0xffff
    where y = 0xffff .&. (fromIntegral x)

compile :: String -> Either String Expression
compile expr = Expression <$> compileToken `mapM` words expr
    where compileToken "+"   = Right Add
          compileToken "-"   = Right Sub
          compileToken "*"   = Right Mul
          compileToken "/"   = Right Div
          compileToken "%"   = Right Mod
          compileToken ">>"  = Right ShRight
          compileToken "<<"  = Right ShLeft
          compileToken "="   = Right Eq
          compileToken ">"   = Right Gt
          compileToken "<"   = Right Lt
          compileToken ">="  = Right Ge
          compileToken "<="  = Right Le
          compileToken "||"  = Right Or
          compileToken "&&"  = Right And
          compileToken "!"   = Right Not
          compileToken "|"   = Right BitOr
          compileToken "&"   = Right BitAnd
          compileToken "^"   = Right BitXor
          compileToken "~"   = Right BitNot
          compileToken "abs" = Right Abs
          compileToken "x"   = Right VarX
          compileToken "y"   = Right VarY
          compileToken "w"   = Right VarW
          compileToken "h"   = Right VarH
          compileToken "t"   = Right VarT
          compileToken "dup" = Right Dup
          compileToken token = case maybeRead token of
                                 Just n -> Right $ Num n
                                 _      -> Left $ "Invalid token: " <> token

evaluate :: ExprMeta -> Expression -> Either String NumT
evaluate meta (Expression items) = go items []
    where
        int :: Bool -> NumT
        int False = 0
        int True  = 1

        bool :: NumT -> Bool
        bool 0 = False
        bool _ = True

        go :: [Item] -> [NumT] -> Either String NumT
        go [] [x]                   = Right $ x
        go [] _                     = Left  $ "Expected one result in stack"

        go (Add:it) (y:x:stack)     = go it (x + y : stack)
        go (Sub:it) (y:x:stack)     = go it (x - y : stack)
        go (Mul:it) (y:x:stack)     = go it (x * y : stack)
        go (Div:it) (y:x:stack)     = go it (x `div` y : stack)
        go (Mod:it) (y:x:stack)     = go it (x `mod` y : stack)
        go (ShRight:it) (y:x:stack) = go it (x `shiftR` fromIntegral y : stack)
        go (ShLeft:it) (y:x:stack)  = go it (x `shiftL` fromIntegral y : stack)
        go (Eq:it) (y:x:stack)      = go it (int (x == y) : stack)
        go (Gt:it) (y:x:stack)      = go it (int (x > y) : stack)
        go (Lt:it) (y:x:stack)      = go it (int (x < y) : stack)
        go (Ge:it) (y:x:stack)      = go it (int (x >= y) : stack)
        go (Le:it) (y:x:stack)      = go it (int (x <= y) : stack)
        go (Or:it) (y:x:stack)      = go it (int ((bool x) || (bool y)) : stack)
        go (And:it) (y:x:stack)     = go it (int ((bool x) && (bool y)) : stack)
        go (Not:it) (x:stack)       = go it (int (not (bool x)) : stack)
        go (BitOr:it) (y:x:stack)   = go it ((x .|. y) : stack)
        go (BitXor:it) (y:x:stack)  = go it ((x `xor` y) : stack)
        go (BitAnd:it) (y:x:stack)  = go it ((x .&. y) : stack)
        go (BitNot:it) (x:stack)    = go it (complement x : stack)
        go (Abs:it) (x:stack)       = go it (abs x : stack)
        go (VarX:it) stack          = go it (getX meta : stack)
        go (VarY:it) stack          = go it (getY meta : stack)
        go (VarW:it) stack          = go it (getW meta : stack)
        go (VarH:it) stack          = go it (getH meta : stack)
        go (Dup:it) (x:stack)       = go it (x : x : stack)
        go ((Num x):it) stack       = go it (x : stack)
        go _ _ = Left "Unexpected number of args passed"

maybeRead :: Read a => String -> Maybe a
maybeRead = fmap fst . listToMaybe . reads

threadPool :: Int -> (a -> IO ()) -> IO (Chan a)
threadPool threads task = do
    input <- newChan
    forM_ [1..threads] $
        \_ -> do
            handle <- dupChan input
            forkOS $ forever $ do
                x <- readChan handle
                task x
    return input

Main.hs

module Main where

import Bitlang

import System.Environment

-- Example: ./Main 1 'x y + abs x y - abs 1 + ^ 2 << 3 % !'
main :: IO ()
main = do
    args <- getArgs
    case args of
      ("1":expr:_)  -> generateImage One expr 
      ("16":expr:_) -> generateImage Sixteen expr 
      _             -> error "Usage: ./Main <depth> <expr>"

generateImage depth str = do
    case compile str of
      Left err   -> error err
      Right expr -> do
          let meta = ImageMeta 500 500
          result <- generate depth expr meta
          case result of
            Right _  -> return ()
            Left err -> error err