From e2bdd7eae231f63e52760daaf5d419d674957ba0 Mon Sep 17 00:00:00 2001
From: Michiel Borkent <michielborkent@gmail.com>
Date: Sat, 11 Apr 2020 21:36:33 +0200
Subject: [PATCH] [#348] nrepl: support multiple top level expressions

---
 src/babashka/impl/nrepl_server.clj       | 55 +++++++++++++-----------
 test/babashka/impl/nrepl_server_test.clj | 10 ++++-
 2 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/src/babashka/impl/nrepl_server.clj b/src/babashka/impl/nrepl_server.clj
index 8b9026ac..db62b227 100644
--- a/src/babashka/impl/nrepl_server.clj
+++ b/src/babashka/impl/nrepl_server.clj
@@ -3,8 +3,10 @@
   (:refer-clojure :exclude [send future binding])
   (:require [babashka.impl.bencode.core :refer [write-bencode read-bencode]]
             [clojure.string :as str]
+            [clojure.tools.reader.reader-types :as r]
             [sci.core :as sci]
-            [sci.impl.interpreter :refer [eval-string*]]
+            [sci.impl.interpreter :refer [eval-string* eval-form]]
+            [sci.impl.parser :as p]
             [sci.impl.utils :as sci-utils]
             [sci.impl.vars :as vars])
   (:import [java.io StringWriter OutputStream InputStream PushbackInputStream EOFException BufferedOutputStream]
@@ -38,30 +40,35 @@
 
 (defn eval-msg [ctx o msg]
   (try
-    (let [ns-str (get msg :ns)
-          sci-ns (when ns-str (sci-utils/namespace-object (:env ctx) (symbol ns-str) true nil))
-          sw (StringWriter.)]
-      (sci/with-bindings (cond-> {sci/out sw}
+    (let [code-str (get msg :code)
+          reader (r/indexing-push-back-reader (r/string-push-back-reader code-str))
+          ns-str (get msg :ns)
+          sci-ns (when ns-str (sci-utils/namespace-object (:env ctx) (symbol ns-str) true nil))]
+      (when @dev? (println "current ns" (vars/current-ns-name)))
+      (sci/with-bindings (cond-> {}
                            sci-ns (assoc vars/current-ns sci-ns))
-        (when @dev? (println "current ns" (vars/current-ns-name)))
-        (let [code-str (get msg :code)
-              value (if (str/blank? code-str)
-                      ::nil
-                      (eval-string* ctx code-str))
-              out-str (not-empty (str sw))
-              env (:env ctx)]
-          (swap! env update-in [:namespaces 'clojure.core]
-                 (fn [core]
-                   (assoc core
-                          '*1 value
-                          '*2 (get core '*1)
-                          '*3 (get core '*2))))
-          (when @dev? (println "out str:" out-str))
-          (when out-str
-            (send o (response-for msg {"out" out-str})))
-          (send o (response-for msg (cond-> {"ns" (vars/current-ns-name)}
-                                      (not (identical? value ::nil)) (assoc "value" (pr-str value)))))
-          (send o (response-for msg {"status" #{"done"}})))))
+        (loop []
+          (let [sw (StringWriter.)
+                form (p/parse-next ctx reader)
+                value (if (identical? :edamame.impl.parser/eof form) ::nil
+                          (sci/with-bindings {sci/out sw}
+                            (eval-form ctx form)))
+                out-str (not-empty (str sw))
+                env (:env ctx)]
+            (swap! env update-in [:namespaces 'clojure.core]
+                   (fn [core]
+                     (assoc core
+                            '*1 value
+                            '*2 (get core '*1)
+                            '*3 (get core '*2))))
+            (when @dev? (println "out str:" out-str))
+            (when out-str
+              (send o (response-for msg {"out" out-str})))
+            (send o (response-for msg (cond-> {"ns" (vars/current-ns-name)}
+                                        (not (identical? value ::nil)) (assoc "value" (pr-str value)))))
+            (when (not (identical? ::nil value))
+              (recur)))))
+      (send o (response-for msg {"status" #{"done"}})))
     (catch Exception ex
       (swap! (:env ctx) update-in [:namespaces 'clojure.core]
              assoc '*e ex)
diff --git a/test/babashka/impl/nrepl_server_test.clj b/test/babashka/impl/nrepl_server_test.clj
index f0947bbc..0892f9d7 100644
--- a/test/babashka/impl/nrepl_server_test.clj
+++ b/test/babashka/impl/nrepl_server_test.clj
@@ -76,7 +76,15 @@
                                        "id" (new-id!)
                                        "ns" "unicorn"})
             (let [reply (read-reply in session @id)]
-              (is (= "unicorn" (:value reply)))))))
+              (is (= "unicorn" (:value reply))))))
+        (testing "multiple top level expressions results in two value replies"
+          (bencode/write-bencode os {"op" "eval"
+                                     "code" "(+ 1 2 3) (+ 1 2 3)"
+                                     "session" session
+                                     "id" (new-id!)})
+          (let [reply-1 (read-reply in session @id)
+                reply-2 (read-reply in session @id)]
+            (is (= "6" (:value reply-1) (:value reply-2))))))
       (testing "load-file"
         (bencode/write-bencode os {"op" "load-file" "file" "(ns foo) (defn foo [] :foo)" "session" session "id" (new-id!)})
         (read-reply in session @id)