From 54bfcaf1ff40570850c40c4af2511bf2443a868d Mon Sep 17 00:00:00 2001
From: Seeker <meaningseeking@protonmail.com>
Date: Sat, 26 Dec 2020 23:48:20 -0800
Subject: Add support for squiggly heredocs

---
 AUTHORS                             |  1 +
 include/mruby/compile.h             |  1 +
 mrbgems/mruby-compiler/core/parse.y | 63 ++++++++++++++++++++++++++++++++++---
 test/t/literals.rb                  | 24 ++++++++++++++
 4 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/AUTHORS b/AUTHORS
index 86428f554..97ab00311 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -45,3 +45,4 @@ of this list.
    RIZAL Reckordp
    Rory O'Connell
    John Bampton
+   Seeker
diff --git a/include/mruby/compile.h b/include/mruby/compile.h
index 9b062e60b..d94e4c1f5 100644
--- a/include/mruby/compile.h
+++ b/include/mruby/compile.h
@@ -99,6 +99,7 @@ enum mrb_string_type {
 /* heredoc structure */
 struct mrb_parser_heredoc_info {
   mrb_bool allow_indent:1;
+  mrb_bool remove_indent:1;
   mrb_bool line_head:1;
   enum mrb_string_type type;
   const char *term;
diff --git a/mrbgems/mruby-compiler/core/parse.y b/mrbgems/mruby-compiler/core/parse.y
index e4a3b9f28..77279b8a8 100644
--- a/mrbgems/mruby-compiler/core/parse.y
+++ b/mrbgems/mruby-compiler/core/parse.y
@@ -1351,6 +1351,55 @@ heredoc_treat_nextline(parser_state *p)
 static void
 heredoc_end(parser_state *p)
 {
+  parser_heredoc_info *info = parsing_heredoc_inf(p);
+  if (info->remove_indent) {
+    mrb_bool counting = TRUE;
+    size_t indent = -1;
+    node *list = info->doc;
+    node *list2 = NULL;
+    while (list) {
+      if (list->car->car == NODE_STR) {
+        node *pair = list->car->cdr;
+        const char *str = (char*) pair->car;
+        size_t len = (size_t) pair->cdr;
+        if (counting) {
+          list2 = push(list2, pair);
+        }
+        size_t spaces = 0;
+        for (size_t i = 0; i < len; i++) {
+          if (counting) {
+            if (ISSPACE(str[i])) {
+              ++spaces;
+            }
+            else {
+              counting = FALSE;
+              if (indent == -1 || spaces < indent) {
+                indent = spaces;
+              }
+            }
+          }
+          if (str[i] == '\n') {
+            counting = TRUE;
+            break;
+          }
+        }
+      }
+      else {
+        counting = FALSE;
+      }
+      list = list->cdr;
+    }
+    if (indent > 0) {
+      while (list2) {
+        node *pair = list2->car;
+        const char *str = (char*) pair->car;
+        size_t len = (size_t) pair->cdr;
+        pair->car = (node*) (str + indent);
+        pair->cdr = (node*) (len - indent);
+        list2 = list2->cdr;
+      }
+    }
+  }
   p->parsing_heredoc = p->parsing_heredoc->cdr;
   if (p->parsing_heredoc == NULL) {
     p->lstate = EXPR_BEG;
@@ -1499,7 +1548,7 @@ heredoc_end(parser_state *p)
 %token tANDDOT            /* &. */
 %token tSYMBEG tREGEXP_BEG tWORDS_BEG tSYMBOLS_BEG
 %token tSTRING_BEG tXSTRING_BEG tSTRING_DVAR tLAMBEG
-%token <nd> tHEREDOC_BEG  /* <<, <<- */
+%token <nd> tHEREDOC_BEG  /* <<, <<-, <<~ */
 %token tHEREDOC_END tLITERAL_DELIM tHD_LITERAL_DELIM
 %token <nd> tHD_STRING_PART tHD_STRING_MID
 
@@ -4906,6 +4955,7 @@ heredoc_identifier(parser_state *p)
   int c;
   int type = str_heredoc;
   mrb_bool indent = FALSE;
+  mrb_bool squiggly = FALSE;
   mrb_bool quote = FALSE;
   node *newnode;
   parser_heredoc_info *info;
@@ -4915,8 +4965,11 @@ heredoc_identifier(parser_state *p)
     pushback(p, c);
     return 0;
   }
-  if (c == '-') {
-    indent = TRUE;
+  if (c == '-' || c == '~') {
+    if (c == '-')
+      indent = TRUE;
+    if (c == '~')
+      squiggly = TRUE;
     c = nextc(p);
   }
   if (c == '\'' || c == '"') {
@@ -4943,6 +4996,7 @@ heredoc_identifier(parser_state *p)
     if (! identchar(c)) {
       pushback(p, c);
       if (indent) pushback(p, '-');
+      if (squiggly) pushback(p, '~');
       return 0;
     }
     newtok(p);
@@ -4959,7 +5013,8 @@ heredoc_identifier(parser_state *p)
   if (! quote)
     type |= STR_FUNC_EXPAND;
   info->type = (string_type)type;
-  info->allow_indent = indent;
+  info->allow_indent = indent || squiggly;
+  info->remove_indent = squiggly;
   info->line_head = TRUE;
   info->doc = NULL;
   p->heredocs_from_nextline = push(p->heredocs_from_nextline, newnode);
diff --git a/test/t/literals.rb b/test/t/literals.rb
index 6344219aa..93652d8a4 100644
--- a/test/t/literals.rb
+++ b/test/t/literals.rb
@@ -162,6 +162,26 @@ qq
 QQ2
       "o")
 
+  r = <<~RRR
+    rrr
+      rrr
+  RRR
+  s = <<~"SSS"
+    sss
+      sss
+  SSS
+  t = <<~'TTT'
+    ttt
+      ttt
+  TTT
+  u = [<<~UUU1   , <<~"UUU2"   , <<~'UUU3' ]
+    u#{1}u
+  UUU1
+    u#{2}u
+  UUU2
+    u#{3}u
+  UUU3
+
   w = %W( 1 #{<<WWW} 3
 www
 WWW
@@ -197,6 +217,10 @@ ZZZ
   assert_equal [1, "nn1\n", 3, 4], n
   assert_equal "a $ q\n $ c $ d", q1
   assert_equal "l $ mqq\nn $ o", q2
+  assert_equal "rrr\n  rrr\n", r
+  assert_equal "sss\n  sss\n", s
+  assert_equal "ttt\n  ttt\n", t
+  assert_equal ["u1u\n", "u2u\n", "u\#{3}u\n"], u
   assert_equal ["1", "www\n", "3", "4", "5"], w
   assert_equal [1, "foo 222 333\n 444\n5\n bar\n6\n", 9], x
   assert_equal "", z
-- 
cgit v1.2.3