From c45767a00591218eba15e499ea30b1312a74e503 Mon Sep 17 00:00:00 2001
From: Colin <colin@uninsane.org>
Date: Tue, 4 Jul 2023 10:02:58 +0000
Subject: [PATCH] handle arguments to the `#!nix-shell` directive

---
 corpus/payload.txt                      |  26 +++++-
 corpus/shebang_with_sub_interpreter.txt | 115 +++++++++++++++++++++++-
 grammar.js                              |  38 +++++++-
 3 files changed, 173 insertions(+), 6 deletions(-)

diff --git a/corpus/payload.txt b/corpus/payload.txt
index 4305fda..eeaf496 100644
--- a/corpus/payload.txt
+++ b/corpus/payload.txt
@@ -11,5 +11,29 @@ echo "hello, world!"
   (first_line)
   (next_line
     (annotation_line
-      (nix_shell_args)))
+      (nix_shell_args
+        (interpreter
+          (sh_arg)))))
+  (next_line))
+
+==================
+nix-shell with interspersed payload
+==================
+
+#!/usr/bin/env nix-shell
+echo hello
+echo again
+#!nix-shell -i bash
+echo everybody
+---
+
+(source_file
+  (first_line)
+  (next_line)
+  (next_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (interpreter
+          (sh_arg)))))
   (next_line))
diff --git a/corpus/shebang_with_sub_interpreter.txt b/corpus/shebang_with_sub_interpreter.txt
index 026c400..978e7ac 100644
--- a/corpus/shebang_with_sub_interpreter.txt
+++ b/corpus/shebang_with_sub_interpreter.txt
@@ -1,5 +1,5 @@
 ==================
-Minimal nix-shell invocation
+minimal nix-shell invocation with interpreter
 ==================
 
 #!/usr/bin/env nix-shell
@@ -10,4 +10,115 @@ Minimal nix-shell invocation
   (first_line)
   (next_line
     (annotation_line
-      (nix_shell_args))))
+      (nix_shell_args
+        (interpreter
+          (sh_arg))))))
+
+==================
+non-interpreter argument
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell -i bash -p gnused
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (interpreter
+          (sh_arg))
+        (sh_arg)
+        (sh_arg)))))
+
+==================
+non-interpreter flags
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell -i bash -p gnused awk
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (interpreter
+          (sh_arg))
+        (sh_arg)
+        (sh_arg)
+        (sh_arg)))))
+
+==================
+quoted arguments
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell -i bash arg\ one 'arg"Two' "'argThree"
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (interpreter
+          (sh_arg))
+        (sh_arg)
+        (sh_arg)
+        (sh_arg)))))
+
+==================
+order-independent
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell -p gnused -i bash
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (sh_arg)
+        (sh_arg)
+        (interpreter
+          (sh_arg))))))
+
+==================
+empty nix-shell
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line)))
+
+==================
+multiple annotation lines
+==================
+
+#!/usr/bin/env nix-shell
+#!nix-shell -p gnused
+#!nix-shell -i bash
+---
+
+(source_file
+  (first_line)
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (sh_arg)
+        (sh_arg))))
+  (next_line
+    (annotation_line
+      (nix_shell_args
+        (interpreter
+          (sh_arg))))))
diff --git a/grammar.js b/grammar.js
index 9192f83..81b1ac2 100644
--- a/grammar.js
+++ b/grammar.js
@@ -9,10 +9,42 @@ module.exports = grammar({
         $.annotation_line,
         repeat(/./),
     )),
-    annotation_line: $ => seq($._shebang_open, $._opt_ws, $._nix_shell, $.nix_shell_args),
+    annotation_line: $ => seq($._shebang_open, $._opt_ws, $._nix_shell, optional($.nix_shell_args)),
 
-    // TODO: actually parse nix-shell arguments
-    nix_shell_args: $ => ' -i bash',
+    // nix acts as if it actually just `exec`'s whichever line contains `#! nix-shell` in the style of `sh`'s `exec`.
+    // so arguments can be quoted, invoke subshells, interpolate environment variables, etc.
+    // the full scope cannot be covered, but the minimal scope here gets 99% of use cases.
+    nix_shell_args: $ => repeat1(seq($._ws, $._nix_shell_arg)),
+    _nix_shell_arg: $ => choice(
+      seq('-i', $._opt_ws, $.interpreter),
+      $.sh_arg,
+    ),
+    interpreter: $ => $.sh_arg,
+    sh_arg: $ => choice(
+      $._sh_lit,
+      $._sh_quote1,
+      $._sh_quote2,
+    ),
+    _sh_lit: $ => repeat1(choice(
+      /[^ \\]/,
+      seq('\\', /./),
+    )),
+    _sh_quote1: $ => seq(
+      '\'',
+      repeat(choice(
+        /[^'\\]/,
+        seq('\\', /./),
+      )),
+      '\'',
+    ),
+    _sh_quote2: $ => seq(
+      '"',
+      repeat(choice(
+        /[^"\\]/,
+        seq('\\', /./),
+      )),
+      '"',
+    ),
 
     _shebang_open: $ => '#!',
     // TODO: allow `usr/bin/env`, `env`