A lazy brainfuck interpreter in Clojure

brainfuck.clj

(ns brainfuck
  "A lazy brainfuck interpreter.

  Memory consists of a fixed-length array of bytes in the range
  0..255, where arithmetic wraps around. Moving the data pointer
  out of the range of memory throws a SEGFAULT error.

  Input must be given as a string or sequence when calling one
  of the eval functions, and if the program tries to read after
  input has been exhausted an exception will be thrown.")

(defn safe-get
  [m k]
  (doto (get m k) (assert)))

(def instruction-types
  {\> :move-data-pointer
   \< :move-data-pointer
   \+ :modify-data
   \- :modify-data
   \. :output
   \, :input
   \[ :loop
   \] :loop})

(defn parse-brackets
  "Returns a map from the address of each bracket to the address of its
  matching bracket."
  ([program] (parse-brackets program 0 () {}))
  ([program i bracket-stack matching-brackets]
     (if-let [[c & cs] (seq program)]
       (case c
         \[ (recur cs (inc i) (conj bracket-stack i) matching-brackets)
         \] (if (empty? bracket-stack)
              (throw (ex-info "Parse error -- unexpected ]" {:char i}))
              (recur cs (inc i) (pop bracket-stack)
                     (let [j (peek bracket-stack)]
                       (assoc matching-brackets i j j i))))
         (recur cs (inc i) bracket-stack matching-brackets))
       (if (empty? bracket-stack)
         matching-brackets
         (throw (Exception. "Parse error -- unmatched bracket!"))))))

(defn init-state
  [mem-size program input-seq]
  (let [program (filterv instruction-types program)]
    {:mem                 (vec (repeat mem-size 0))
     :program             program
     :matching-brackets   (parse-brackets program)
     :instruction-pointer 0
     :data-pointer        0
     :input               input-seq}))

(defn current-instruction
  "Returns nil if the instruction pointer has been incremented past the
  end of the program."
  [state]
  (let [{:keys [program instruction-pointer]} state]
    (get program instruction-pointer)))

(defn inc-instruction-pointer
  [state]
  (update-in state [:instruction-pointer] inc))

(defmulti step
  "Returns [new-state output-byte] where output-byte is nil if nothing
  was output. Will return a state with :halted? set to true if the
  program is finished."
  #(instruction-types (current-instruction %)))

(defn no-output [state] [state nil])

(defmethod step nil [state]
  (-> state
      (assoc :halted? true)
      (no-output)))

(defmethod step :move-data-pointer
  [{:keys [data-pointer mem] :as state}]
  (let [inc-or-dec ({\> inc \< dec} (current-instruction state))
        data-pointer' (inc-or-dec data-pointer)]
    (if (< -1 data-pointer' (count mem))
      (-> state
          (inc-instruction-pointer)
          (assoc :data-pointer data-pointer')
          (no-output))
      (throw (ex-info "SEGFAULT" {:state state})))))

(defmethod step :modify-data
  [{:keys [mem data-pointer] :as state}]
  (let [byte (safe-get mem data-pointer)
        inc-or-dec ({\+ inc \- dec} (current-instruction state))
        byte' (mod (inc-or-dec byte) 256)]
    (-> state
        (inc-instruction-pointer)
        (assoc-in [:mem data-pointer] byte')
        (no-output))))

(defmethod step :output
  [{:keys [mem data-pointer] :as state}]
  (let [byte (safe-get mem data-pointer)]
    [(-> state (inc-instruction-pointer)) byte]))

(defmethod step :input
  [{:keys [data-pointer input] :as state}]
  (if-let [byte (first input)]
    (do (assert (<= 0 byte 255))
        (-> state
            (assoc-in [:mem data-pointer] byte)
            (assoc :input (rest input))
            (inc-instruction-pointer)
            (no-output)))
    (throw (ex-info "No more input!" {:state state}))))

(defmethod step :loop
  [{:keys [mem data-pointer instruction-pointer matching-brackets] :as state}]
  (let [jump-test-fn ({\[ zero? \] pos?} (current-instruction state))]
    (if (jump-test-fn (safe-get mem data-pointer))
      (let [instruction-pointer' (inc (safe-get matching-brackets instruction-pointer))]
        (-> state
            (assoc :instruction-pointer instruction-pointer')
            (no-output)))
      (-> state
          (inc-instruction-pointer)
          (no-output)))))

(defn states
  "Returns a lazy seq of program states, starting with the input state."
  [state]
  (->> (iterate (comp step first) [state nil])
       (map first)
       (take-while (complement :halted?))))

(defn output
  "Returns a lazy seq of output bytes."
  [state]
  (->> (iterate (comp step first) [state nil])
       (take-while (comp (complement :halted?) first))
       (keep second)))

(defn eval-bytes
  "Returns a lazy output-seq. Both input-seq and
  output-seq are seqs of integers in the range 0..255."
  [mem-size program input-seq]
  (output (init-state mem-size program input-seq)))

(defn eval-chars
  "Returns a lazy output-seq. Both input-seq and
  output-seq are seqs of Characters such that (int c)
  is in the range 0..255."
  [mem-size program input-seq]
  (->> input-seq
       (map int)
       (eval-bytes mem-size program)
       (map char)))

(defn eval-string
  "Returns a string; input is same as eval-chars since strings can be
  treated as seqs of chars."
  [mem-size program s]
  (apply str (eval-chars mem-size program s)))

(comment

  ;; Hello World, straight from the wikipedia
  (def hello-world "++++++++[>++++[>++>+++>+++>+<<<<-]>+>+>->>+[<]<-]>>.>---.+++++++..+++.>>.<-.<.+++.------.--------.>>+.>++.")
  (eval-string 20 hello-world "")
  ;; => "Hello World!\n"


  ;; Slightly modified from the wikipedia version, detects EOF as a null byte,
  ;; and correspondingly prints a null byte at the end.
  (def ROT13
    "-,[+                         Read first character and start outer character reading loop
         -[                       Skip forward if character is 0
             >>++++[>++++++++<-]  Set up divisor (32) for division loop
                                    (MEMORY LAYOUT: dividend copy remainder divisor quotient zero zero)
             <+<-[                Set up dividend (x minus 1) and enter division loop
                 >+>+>-[>>>]      Increase copy and remainder / reduce divisor / Normal case: skip forward
                 <[[>+<-]>>+>]    Special case: move remainder back to divisor and increase quotient
                 <<<<<-           Decrement dividend
             ]                    End division loop
         ]>>>[-]+                 End skip loop; zero former divisor and reuse space for a flag
         >--[-[<->+++[-]]]<[         Zero that flag unless quotient was 2 or 3; zero quotient; check flag
             ++++++++++++<[       If flag then set up divisor (13) for second division loop
                                    (MEMORY LAYOUT: zero copy dividend divisor remainder quotient zero zero)
                 >-[>+>>]         Reduce divisor; Normal case: increase remainder
                 >[+[<+>-]>+>>]   Special case: increase remainder / move it back to divisor / increase quotient
                 <<<<<-           Decrease dividend
             ]                    End division loop
             >>[<+>-]             Add remainder back to divisor to get a useful 13
             >[                   Skip forward if quotient was 0
                 -[               Decrement quotient and skip forward if quotient was 1
                     -<<[-]>>     Zero quotient and divisor if quotient was 2
                 ]<<[<<->>-]>>    Zero divisor and subtract 13 from copy if quotient was 1
             ]<<[<<+>>-]          Zero divisor and add 13 to copy if quotient was 0
         ]                        End outer skip loop (jump to here if ((character minus 1)/32) was not 2 or 3)
         <[-]                     Clear remainder from first division if second division was skipped
         <.[-]                    Output ROT13ed character from copy and clear it
         <-,                      Read next character
     ]                            End character reading loop
     .                            Print terminal null byte")

  (eval-string 20 ROT13 "SEKRET MSG\0")
  ;; => "FRXERG ZFT"

  (->> "SEKRET MSG\0"
       (eval-string 20 ROT13)
       (eval-string 20 ROT13))
  ;; => "SEKRET MSG"

  )