Add variable expansion inside of heredocs

corpus/commands.txt

@@ -124,6 +124,26 @@ JS
     (argument)
     (heredoc_redirect (heredoc)))))
 
+===============================
+Heredocs with variables
+===============================
+
+node <<JS
+a $B ${C}
+JS
+
+exit
+
+---
+
+(program
+  (command (simple_command
+    (command_name)
+    (heredoc_redirect (heredoc
+      (expansion (variable_name))
+      (operator_expansion (variable_name))))))
+  (command (simple_command (command_name))))
+
 ===============================
 Variable expansions
 ===============================

grammar.js

@@ -4,7 +4,10 @@ module.exports = grammar({
   inline: $ => [$.control_operator],
 
   externals: $ => [
-    $.heredoc
+    $._simple_heredoc,
+    $._heredoc_beginning,
+    $._heredoc_middle,
+    $._heredoc_end
   ],
 
   rules: {
@@ -93,6 +96,19 @@ module.exports = grammar({
       $.heredoc
     ),
 
+    heredoc: $ => choice(
+      $._simple_heredoc,
+      seq(
+        $._heredoc_beginning,
+        repeat(choice(
+          $.expansion,
+          $.operator_expansion,
+          $._heredoc_middle
+        )),
+        $._heredoc_end
+      )
+    ),
+
     file_descriptor: $ => token(prec(1, /\d+/)),
 
     leading_word: $ => /[^\s=|;:{}]+/,

src/grammar.json

@@ -380,6 +380,48 @@
         }
       ]
     },
+    "heredoc": {
+      "type": "CHOICE",
+      "members": [
+        {
+          "type": "SYMBOL",
+          "name": "_simple_heredoc"
+        },
+        {
+          "type": "SEQ",
+          "members": [
+            {
+              "type": "SYMBOL",
+              "name": "_heredoc_beginning"
+            },
+            {
+              "type": "REPEAT",
+              "content": {
+                "type": "CHOICE",
+                "members": [
+                  {
+                    "type": "SYMBOL",
+                    "name": "expansion"
+                  },
+                  {
+                    "type": "SYMBOL",
+                    "name": "operator_expansion"
+                  },
+                  {
+                    "type": "SYMBOL",
+                    "name": "_heredoc_middle"
+                  }
+                ]
+              }
+            },
+            {
+              "type": "SYMBOL",
+              "name": "_heredoc_end"
+            }
+          ]
+        }
+      ]
+    },
     "file_descriptor": {
       "type": "TOKEN",
       "content": {
@@ -423,7 +465,19 @@
   "externals": [
     {
       "type": "SYMBOL",
-      "name": "heredoc"
+      "name": "_simple_heredoc"
+    },
+    {
+      "type": "SYMBOL",
+      "name": "_heredoc_beginning"
+    },
+    {
+      "type": "SYMBOL",
+      "name": "_heredoc_middle"
+    },
+    {
+      "type": "SYMBOL",
+      "name": "_heredoc_end"
     }
   ],
   "inline": [

src/scanner.cc

@@ -7,7 +7,10 @@ namespace {
 using std::wstring;
 
 enum TokenType {
-  HEREDOC
+  SIMPLE_HEREDOC,
+  HEREDOC_BEGINNING,
+  HEREDOC_MIDDLE,
+  HEREDOC_END,
 };
 
 struct Scanner {
@@ -25,32 +28,73 @@ struct Scanner {
 
   void deserialize(TSExternalTokenState state) {}
 
-  bool scan(TSLexer *lexer, const bool *valid_symbols) {
-    wstring heredoc_content;
-
+  bool scan_heredoc_end_identifier(TSLexer *lexer) {
+    current_leading_word.clear();
     while (iswalpha(lexer->lookahead)) {
-      heredoc_content += lexer->lookahead;
+      current_leading_word += lexer->lookahead;
+      advance(lexer);
+    }
+    return current_leading_word == heredoc_identifier;
+  }
+
+  bool scan_heredoc_content(TSLexer *lexer, TokenType middle_type, TokenType end_type) {
+    bool did_advance = false;
+
+    for (;;) {
+      switch (lexer->lookahead) {
+        case '\0': {
+          lexer->result_symbol = end_type;
+          return true;
+        }
+
+        case '$': {
+          lexer->result_symbol = middle_type;
+          return did_advance;
+        }
+
+        case '\n': {
+          did_advance = true;
+          advance(lexer);
+          if (scan_heredoc_end_identifier(lexer)) {
+            lexer->result_symbol = end_type;
+            return true;
+          }
+          break;
+        }
+
+        default: {
+          did_advance = true;
+          advance(lexer);
+          break;
+        }
+      }
+    }
+  }
+
+  bool scan(TSLexer *lexer, const bool *valid_symbols) {
+    if (valid_symbols[HEREDOC_MIDDLE]) {
+      return scan_heredoc_content(lexer, HEREDOC_MIDDLE, HEREDOC_END);
+    }
+
+    heredoc_identifier.clear();
+    while (iswalpha(lexer->lookahead)) {
+      heredoc_identifier += lexer->lookahead;
       advance(lexer);
     }
 
     if (lexer->lookahead != '\n') return false;
+    advance(lexer);
 
-    wstring leading_word;
-
-    for (;;) {
-      advance(lexer);
-
-      while (iswalpha(lexer->lookahead)) {
-        leading_word += lexer->lookahead;
-        advance(lexer);
-      }
-
-      if (leading_word == heredoc_content || lexer->lookahead == '\0') break;
-      if (lexer->lookahead == '\n') leading_word.clear();
+    if (scan_heredoc_end_identifier(lexer)) {
+      lexer->result_symbol = SIMPLE_HEREDOC;
+      return true;
     }
 
-    return true;
+    return scan_heredoc_content(lexer, HEREDOC_BEGINNING, SIMPLE_HEREDOC);
   }
+
+  wstring heredoc_identifier;
+  wstring current_leading_word;
 };
 
 }