bf.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

module Main where {
import Data.Maybe (catMaybes);
import Data.Function (on);
import Data.List (group,intersperse,sort,sortBy,groupBy);
import qualified Control.Monad.State as S;
--import Control.Monad.Writer;
import Text.ParserCombinators.ReadP;

data Code
    = PutChar Int -- putchar(p[n]);
    | GetChar Int -- p[n] = getchar();
    | Advance Int -- p += n;
    | Increment Int Int -- p[i] += a;
    | While [Code] -- while(*p) { ... } 
    | Set Int Int -- p[i] = a;
    | AddTo [(Int,Int)] (Int,Int)
--    deriving (Eq,Ord,Show);
    deriving (Eq,Ord);

wsLF :: ShowS;
wsLF = showString "\n";
wsTab :: ShowS;
wsTab = showString "\t";
wsSpace :: ShowS;
wsSpace = showString " ";

wsInt :: Int -> ShowS;
wsInt n
    | n > 0 = wsSpace . bits n . wsLF
    | n < 0 = wsTab . bits (-n) . wsLF
    | n == 0 = wsSpace . wsSpace . wsLF
where {
    bits 0 = id;
    bits x = bits (x `div` 2)
    . case (x `mod` 2) of {
        0 -> wsSpace;
        1 -> wsTab;
    };
};


-- The top of the stack at the beginning of a Brainfuck command is always the index of the pointer.
data WhitespaceBytecode
    = PushInt Int -- Push n
    | StackDiscard -- Pop the top item on the stack.
    | StackDup -- Duplicate the top item on the stack.
    | StackSwap -- Swap the top two items on the stack.
    | InputChar -- Input a character and put it on the top of the stack.
    | OutputChar -- Output top of stack.
    | StackAdd -- Pop and add the top two items on the stack and push the sum on the stack.
    | StackMul -- Pop and multiply the top two items on the stack and push the product on the stack.
    | StackMod -- Pop and modulo the top two items on the stack and push the product on the stack.
    | HeapRetrieve -- Pop the top item of the stack and use it as an index to retrieve an item from the heap. Push the heap item.
    | HeapStore -- Pop the value and then the item from the stack. Store the value on the heap at the index.
    | MarkLocation Int -- Mark a location. This does not alter the stack.
    | JumpAlways Int -- Jump to the location unconditionally. This does not alter the stack.
    | JumpZero Int -- Pop a value from the stack and jump to the given location if it is zero.
    | JumpNegative Int -- Pop a value from the stack and jump to the given location if it is negative.
    | WhitespaceTerminate
    deriving (Show)
    ;

-- Print Whitespace bytecode as whitespace.
whitespacePrint :: WhitespaceBytecode -> ShowS;
whitespacePrint (PushInt i)
    = wsSpace
    . wsSpace
    . wsInt i
    ;
whitespacePrint (StackDiscard)
    = wsSpace
    . wsLF . wsLF
    ;
whitespacePrint (StackDup)
    = wsSpace
    . wsLF . wsSpace
    ;
whitespacePrint (StackSwap)
    = wsSpace
    . wsLF . wsTab
    ;
whitespacePrint (InputChar)
    = wsTab . wsLF
    . wsTab . wsSpace
    ;
whitespacePrint (OutputChar)
    = wsTab . wsLF
    . wsSpace . wsSpace
    ;
whitespacePrint (StackAdd)
    = wsTab . wsSpace
    . wsSpace . wsSpace
    ;
whitespacePrint (StackMul)
    = wsTab . wsSpace
    . wsSpace . wsLF
    ;
whitespacePrint (StackMod)
    = wsTab . wsSpace
    . wsTab . wsTab
    ;
whitespacePrint (HeapRetrieve)
    = wsTab . wsTab
    . wsTab
    ;
whitespacePrint (HeapStore)
    = wsTab . wsTab
    . wsSpace
    ;
whitespacePrint (MarkLocation loc)
    = wsLF
    . wsSpace . wsSpace
    . wsInt loc
    ;
whitespacePrint (JumpAlways loc)
    = wsLF
    . wsSpace . wsLF
    . wsInt loc
    ;
whitespacePrint (JumpZero loc)
    = wsLF
    . wsTab . wsSpace
    . wsInt loc
    ;
whitespacePrint (JumpNegative loc)
    = wsLF
    . wsTab . wsTab
    . wsInt loc
    ;
whitespacePrint (WhitespaceTerminate)
    = wsLF
    . wsLF . wsLF
    ;

-- With the base index at the top of the stack, push the sum of the base pointer and the index.
whitecodeRelativeIndex i =
    [StackDup -- [base,base]
    ,PushInt i -- [base,base,i]
    ,StackAdd -- [base,base+i]
    ];

-- Translate Brainfuck bytecode to Whitespace bytecode.
whitespaceBytecode :: Code -> S.State Int [WhitespaceBytecode];
whitespaceBytecode (PutChar i) = return $
    whitecodeRelativeIndex i ++ -- [base,base+i]
    [HeapRetrieve -- [base,array[base+i]]
    ,OutputChar -- [base]
    ];
whitespaceBytecode (GetChar i) = return $
    whitecodeRelativeIndex i ++ -- [base,base+i]
    [InputChar -- [base,base+i,c]
    ,PushInt 256 -- [base,base+i,c,256]
    ,StackMod -- [base,base+i,c % 256]
    ,HeapStore -- [base]
    ];
whitespaceBytecode (Advance i) = return $
    [PushInt i -- [base,i]
    ,StackAdd -- [base+i]
    ];
whitespaceBytecode (Increment i a) = return $
    [StackDup] ++ -- [base,base]
    whitecodeRelativeIndex i ++ -- [base,base,base+i]
    [HeapRetrieve -- [base,base,array[base+i]]
    ,PushInt a -- [base,base,array[base+i],a]
    ,StackAdd -- [base,base,array[base+i]+a]
    ,StackSwap] ++ -- [base,array[base+i]+a,base]
    whitecodeRelativeIndex i ++ -- [base,array[base+i]+a,base,base+i]
    [StackSwap -- [base,array[base+i]+a,base+i,base]
    ,StackDiscard -- [base,array[base+i]+a,base+i]
    ,StackSwap -- [base,base+i,array[base+i]+a]
    ,PushInt 256 -- [base,base+i,array[base+i]+a,256]
    ,StackMod -- [base,base+i,(array[base+i]+a) % 256]
    ,HeapStore -- [base]
    ];
whitespaceBytecode (While xs) = do {
    locBegin <- S.get;
    S.put (locBegin+1);
    middle <- mapM whitespaceBytecode xs;
    locEnd <- S.get;
    S.put (locEnd+1);

    return $
    [MarkLocation locBegin -- [base]
    ,StackDup -- [base,base]
    ,HeapRetrieve -- [base,array[base]]
    ,JumpZero locEnd] ++ -- [base]
    (concat middle) ++
    [JumpAlways locBegin -- [base]
    ,MarkLocation locEnd -- [base]
    ];
};
whitespaceBytecode (Set i a) = return $
    whitecodeRelativeIndex i ++ -- [base,base+i]
    [PushInt a -- [base,base+i,a]
    ,PushInt 256 -- [base,base+i,array[base+i]+a,256]
    ,StackMod -- [base,base+i,(array[base+i]+a) % 256]
    ,HeapStore -- [base]
    ];
whitespaceBytecode (AddTo incs (index,1)) = let {
    -- Map (i,amount) to the command p[i] += p[index]*amount;
    addmul :: (Int,Int) -> [WhitespaceBytecode];
    addmul (i,amount) =
        [StackDup] ++  -- [base,base]
        whitecodeRelativeIndex i ++ -- [base,base,base+i]
        [StackSwap] ++ -- [base,base+i,base]
        whitecodeRelativeIndex index ++ -- [base,base+i,base,base+index]
        [HeapRetrieve -- [base,base+i,base,array[base+index]]
        ,PushInt amount -- [base,base+i,base,array[base+index],amount]
        ,StackMul -- [base,base+i,base,array[base+index]*amount]
        ,StackSwap] ++ -- [base,base+i,array[base+index]*amount,base]
        whitecodeRelativeIndex i ++ -- [base,base+i,array[base+index]*amount,base,base+i]
        [StackSwap -- [base,base+i,array[base+index]*amount,base+i,base]
        ,StackDiscard -- [base,base+i,array[base+index]*amount,base+i]
        ,HeapRetrieve -- [base,base+i,array[base+index]*amount,array[base+i]]
        ,StackAdd -- [base,base+i,array[base+index]*amount+array[base+i]]
        ,PushInt 256 -- [base,base+i,array[base+i]+a,256]
        ,StackMod -- [base,base+i,(array[base+i]+a) % 256]
        ,HeapStore] -- [base]
    ;
} in return $  
    (concatMap addmul incs) ++
    whitecodeRelativeIndex index ++ -- [base,base+index]
    [PushInt 0 -- [base,base+index,0]
    ,HeapStore -- [base]
    ];

-- Print Brainfuck bytecode as C.
instance Show Code where {
    showList xs = foldl (.) id $ intersperse (showString "\n") $ map shows xs;
    showsPrec _ (PutChar 0) = showString "putchar(*p);";
    showsPrec _ (PutChar i) = showString "putchar(p[" . shows i . showString "]);";
    showsPrec _ (GetChar 0) = showString "*p = getchar();";
    showsPrec _ (GetChar i) = showString "p[" . shows i . showString "] = getchar();";
    showsPrec _ (Advance n) = showString "p += " . shows n . showString ";";
    showsPrec _ (Increment 0 n) = showString "*p += " . shows n . showString ";";
    showsPrec _ (Increment i n) = showString "p[" . shows i . showString "] += " . shows n . showString ";";
    showsPrec _ (Set 0 n) = showString "*p = " . shows n . showString ";";
    showsPrec _ (Set i n) = showString "p[" . shows i . showString "] = " . shows n . showString ";";
    showsPrec _ (While xs) = showString "while(*p) {\n" . showList xs . showString "\n}";
    showsPrec _ (AddTo xs (ib,1)) = let {
        makeAdd (index,increment) =
           showString "p[" . shows index . showString "] += " . showMultiply increment  
        ;

        showMultiply 1 = showString "p[" . shows ib . showString "]";
        showMultiply increment = shows increment . showString " * p[" . shows ib . showString "]";
    } in (foldl (.) id $ intersperse (showString "\n") $ map (\x -> makeAdd x . showString ";") xs)
        . showString "p[" . shows ib . showString "] = 0;";
{-
    showsPrec _ (AddTo xs a) = let {
        makeAdd amount (index,increment) =
           showString "p[" . shows index . showString "] += " . shows increment . showString " * " . shows amount;
    } in foldl (.) id $ intersperse (showString "\n") $ map (\x -> makeAdd a x . showString ";") xs;
-}
};

-- Stationary instructions do not move the index.
isStationary :: Code -> Bool;
isStationary (Advance n) = n == 0;
isStationary (While xs) = all isStationary xs;
isStationary _ = True;

-- Parse Brainfuck source code.
bf :: ReadP [Code];
bf = let {
    f t x = do {
        char t;
        return x;
    };
    tree = do {
        tree <- bf';
        return $ While tree;
    };
    bf' = do {
        xs <- many
        (
            (between (char '[') (char ']') tree)
            <++ 
            (choice [
                f ',' $ GetChar 0,
                f '.' $ PutChar 0,
                f '<' $ Advance (-1),
                f '>' $ Advance (1),
                f '-' $ Increment 0 (-1),
                f '+' $ Increment 0 (1)
            ])
        );
        return xs;
    }
} in do {
    xs <- bf';
    eof;
    return xs;
};

-- Optimize Brainfuck bytecode.
optimize :: [Code] -> [Code];
optimize (Advance 0:xs) = optimize $ xs;
optimize (Increment _ 0:xs) = optimize $ xs;

optimize (Advance ia:Increment ib b:xs) = optimize $ Increment (ia+ib) b:Advance ia:xs;
optimize (Advance ia:Set ib b:xs) = optimize $ Set (ia+ib) b:Advance ia:xs;
optimize (Advance ia:PutChar ib:xs) = optimize $ PutChar (ia+ib):Advance ia:xs;
optimize (Advance ia:GetChar ib:xs) = optimize $ GetChar (ia+ib):Advance ia:xs;
optimize (Advance ia:AddTo items (ib,amount) :xs) =
    (AddTo (map (\(i,increment) -> (ia+i,increment)) items) (ib+ia,amount)) : optimize (Advance ia:xs);

optimize (AddTo items (ib,amount) :xs) = let {
    ts = map (\x -> (fst $ head x, sum $ map snd x)) $ groupBy ((==) `on` fst) $ sortBy (compare `on` snd) items;
} in AddTo ts (ib,amount):optimize xs;

optimize (Advance a:Advance b:xs) = optimize $ Advance (a+b):xs;
optimize (Increment ia a:Increment ib b:xs) = 
    if ia == ib
        then optimize $ Increment ia (a+b):xs
        else Increment ia a:optimize (Increment ib b:xs)
    ;
optimize (Set ia a:Increment ib b:xs) =
    if ia == ib
        then optimize $ Set ia (a+b):xs;
        else Set ia a:optimize (Increment ib b:xs);
optimize (While [Increment 0 n]:xs) = 
    if n == 0
        then While []:optimize xs
        else optimize (Set 0 0:xs)
    ;
optimize (While ws:xs) = let {
    isAssignment (Increment _ _) = True;
--    isAssignment (Set _ _) = True;
--    isAssignment (While xs) = all isAssignment xs;
    isAssignment _ = False;
    rec xs = let {
        findIncrement (Increment 0 _) = True;
        findIncrement _ = False;
        Increment _ amount:_ = dropWhile (not . findIncrement) xs;
        getIndex (Increment i s) = Just (i,s);
        getIndex _ = Nothing;
        indexes = catMaybes $ map getIndex xs;
    } in if amount == (-1)
        then AddTo (filter ((/=0) . fst) indexes) (0,-amount)
        else While $ optimize xs
        ;
} in if (all isAssignment ws) && (not $ null ws)
    then (rec ws): optimize xs
    else (While $ optimize ws) : optimize xs;
    ;

optimize (x:xs) = x:optimize xs;
optimize [] = [];

-- Compute a fixed point. f (fixedpoint f x) == fixedpoint f x
fixedPoint :: Eq a => (a -> a) -> a -> a;
fixedPoint f x =
    let {
        y = f x;
    } in if y == x then x else fixedPoint f y;

-- Write Brainfuck bytecode as C.
outputC :: [Code] -> ShowS;
outputC codes = let {
    prologue = showString "int main() { char array[65536] = {0}; char * p = array;";
    epilogue = showString "return 0; }";
} in
      prologue
    . showList codes
    . epilogue
    ;

-- Write Brainfuck bytecode as Whitespace.
outputWhitespace :: [Code] -> ShowS;
outputWhitespace codes = let {
    --codes = fst $ head $ readP_to_S bf $ filter (`elem` "[]<>+-.,") $ str;
    initializeHeap heapSize = do {
        beginLoc <- S.get;
        S.put $ beginLoc + 1;
        return 
            [PushInt 0 -- [base,i]
            ,MarkLocation beginLoc -- [base,i]
            ,StackDup -- [base,i,i]
            ,PushInt 0 -- [base,i,i,0]
            ,HeapStore -- [base,i]
            ,PushInt 1 -- [base,i,1]
            ,StackAdd -- [base,i+1]
            ,StackDup -- [base,i+1,i+1]
            ,PushInt (-heapSize) -- [base,i+1,i+1,-heapSize]
            ,StackAdd -- [base,i+1,i+1-heapSize]
            ,JumpNegative beginLoc -- [base,i+1],
            ,StackDiscard -- [base]
            ];
    };
    heapSize = 6000; -- Should be 30000
    ws :: [WhitespaceBytecode];
    ws = fst $ (flip S.runState) 0 $ do {
        heapInit <- initializeHeap heapSize;
        whitespaceCodes <- mapM whitespaceBytecode codes;
        return $ concat $ ([PushInt 0]:[heapInit] ++ whitespaceCodes ++ [[WhitespaceTerminate]]);
    };
} in foldl (.) id $ map whitespacePrint $ ws;
--} in foldl (.) id $ intersperse (showString "\n") $ map shows  $ ws;

main = do {
    str <- getContents;
    let { codes = fixedPoint optimize $ fst $ head $ readP_to_S bf $ filter (`elem` "[]<>+-.,") $ str; };

    putStrLn $ outputWhitespace codes "";
    --putStrLn $ outputC codes "";
};

}