From 962a3f45787687ec5862977f677758b84ccb0548 Mon Sep 17 00:00:00 2001
From: Eric Wang <eric@wrk.local>
Date: Fri, 5 Nov 2021 17:53:54 +0800
Subject: [PATCH] rustic-babel: support wrap BODY in a fn main if none exists

close #267

related: #266
---
 README.md                 | 38 ++++++++++++++++++++++++++++++++++++--
 rustic-babel.el           | 10 +++++++++-
 test/rustic-babel-test.el | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index a8a3f383..5d79d1df 100644
--- a/README.md
+++ b/README.md
@@ -42,6 +42,7 @@
             - [:features](#features)
             - [:paths](#paths)
             - [:toolchain](#toolchain)
+            - [:main](#main)
     - [Spinner](#spinner)
     - [inline-documentation](#inline-documentation)
         - [Prequisites](#prequisites)
@@ -79,7 +80,7 @@ The other files provide functionality that is similar to some of the features
 of rustic, however can be considered light-weight compared to some rustic's
 functionality.
 
-The shared functions and options exist as aliases in the rust-mode and 
+The shared functions and options exist as aliases in the rust-mode and
 rustic namespace for backwards compatability reasons(rustic has been a fork).
 
 ## Known issues
@@ -499,7 +500,7 @@ then a string, instead of a list, will also be accepted:
 ```
 #+BEGIN_SRC rust :crates '((tokio . 1.0)) :features '((tokio . ("rt-multi-thread" "time")))
   extern crate tokio;
-  
+
   fn main() {
       tokio::runtime::Runtime::new()
           .unwrap()
@@ -541,6 +542,39 @@ fn main() {
 : a,b,c
 ```
 
+#### :main
+
+Auto wrap whole block body in a "fn main" function call if none exists.
+
+This is very useful in most code snippets, so the default value is "yes".
+"no" if you don't want this feature(for example if you don't want regex search slow things down).
+
+```
+#+begin_src rust
+let x = vec![1, 2, 3].iter().map(|&x| x + 1).collect::<Vec<_>>();
+println!("{:?}", x);
+#+end_src
+
+#+results:
+: [2, 3, 4]
+```
+
+NOTE: the regex match would unable to /understand/ `fn main()` within a raw string.
+
+```
+#+begin_src rust :exports both :main yes
+// won't work
+let r = "
+fn main() {
+  let x = "within raw";
+}
+";
+
+// but this works
+printlnt!("fn main{{}}")
+#+end_src
+```
+
 
 ## Spinner
 
diff --git a/rustic-babel.el b/rustic-babel.el
index f9a5cd33..8fdd2053 100644
--- a/rustic-babel.el
+++ b/rustic-babel.el
@@ -259,6 +259,12 @@ directory DIR."
           (insert s)
           (insert dependencies))))))
 
+(defun rustic-babel-ensure-main-wrap (body)
+  "Wrap BODY in a \"fn main\" function call if none exists."
+  (if (string-match "^[ \t]*[fn]+[ \t\n\r]*main[ \t]*(.*)" body)
+      body
+    (format "fn main() {\n%s\n}\n" body)))
+
 (defun org-babel-execute:rustic (body params)
   "Execute a block of Rust code with org-babel.
 
@@ -272,6 +278,7 @@ kill the running process."
       (let* ((default-directory org-babel-temporary-directory)
              (project (rustic-babel-project))
              (dir (setq rustic-babel-dir (expand-file-name project)))
+             (main-p (not (string= (cdr (assq :main params)) "no")))
              (main (expand-file-name "main.rs" (concat dir "/src"))))
         (make-directory (file-name-directory main) t)
         (rustic-babel-cargo-toml dir params)
@@ -286,7 +293,8 @@ kill the running process."
         (let ((default-directory dir)
               (toolchain (cdr (assq :toolchain params))))
           (write-region
-           (concat "#![allow(non_snake_case)]\n" body) nil main nil 0)
+           (concat "#![allow(non_snake_case)]\n"
+                   (if main-p (rustic-babel-ensure-main-wrap body) body)) nil main nil 0)
           (rustic-babel-eval dir toolchain)
           (setq rustic-babel-src-location
                 (set-marker (make-marker) (point) (current-buffer)))
diff --git a/test/rustic-babel-test.el b/test/rustic-babel-test.el
index 79b24acf..88968fa3 100644
--- a/test/rustic-babel-test.el
+++ b/test/rustic-babel-test.el
@@ -187,3 +187,40 @@
     (with-current-buffer buf
       (rustic-test-babel-execute-block buf)
       (should (eq (rustic-test-babel-check-results buf) nil)))))
+
+(ert-deftest rustic-test-babel-ensure-main-wrap()
+  (let* ((string "let x = \"fn main(){}\";")
+         (buf (rustic-test-get-babel-block string)))
+    (with-current-buffer buf
+      (rustic-test-babel-execute-block buf)
+      (should (eq (rustic-test-babel-check-results buf) nil))))
+  (let* ((string "let x = \"fn main(){}\";")
+         (params ":main no")
+         (buf (rustic-test-get-babel-block string params)))
+    (with-current-buffer buf
+      (rustic-test-babel-execute-block buf)
+      (let ((re (format "error: Could not compile `%s`.\n"
+                        (car (reverse (split-string rustic-babel-dir "/"))))))
+        (should (string= re (rustic-test-babel-check-results buf)))))))
+
+(ert-deftest rustic-test-babel-ensure-main-wrap-yes-with-main()
+  (let* ((string "
+                fn main() {
+let x = \"rustic\";
+                  }")
+         (params ":main yes")
+         (buf (rustic-test-get-babel-block string params)))
+    (with-current-buffer buf
+      (rustic-test-babel-execute-block buf)
+      (should (eq (rustic-test-babel-check-results buf) nil)))))
+
+(ert-deftest rustic-test-babel-ensure-main-wrap-no()
+  (let* ((string "
+                fn main() {
+let x = \"rustic\";
+                  }")
+         (params ":main no")
+         (buf (rustic-test-get-babel-block string params)))
+    (with-current-buffer buf
+      (rustic-test-babel-execute-block buf)
+      (should (eq (rustic-test-babel-check-results buf) nil)))))