Shellbang support with flakes

Enables shebang usage of nix shell. All arguments with `#! nix` get
added to the nix invocation. This implementation does NOT set any
additional arguments other than placing the script path itself as the
first argument such that the interpreter can utilize it.

Example below:

```
    #!/usr/bin/env nix
    #! nix shell --quiet
    #! nix nixpkgs#bash
    #! nix nixpkgs#shellcheck
    #! nix nixpkgs#hello
    #! nix --ignore-environment --command bash
    # shellcheck shell=bash
    set -eu
    shellcheck "$0" || exit 1
    function main {
        hello
        echo 0:"$0" 1:"$1" 2:"$2"
    }
    "$@"
```

src/libutil/args.cc

@@ -1,6 +1,7 @@
 #include "args.hh"
 #include "hash.hh"
 
+#include <regex>
 #include <glob.h>
 
 #include <nlohmann/json.hpp>
@@ -54,6 +55,12 @@ std::optional<std::string> needsCompletion(std::string_view s)
 }
 
 void Args::parseCmdline(const Strings & _cmdline)
+{
+    // Default via 5.1.2.2.1 in C standard
+    Args::parseCmdline("", _cmdline);
+}
+
+void Args::parseCmdline(const string & programName, const Strings & _cmdline)
 {
     Strings pendingArgs;
     bool dashDash = false;
@@ -69,6 +76,39 @@ void Args::parseCmdline(const Strings & _cmdline)
     }
 
     bool argsSeen = false;
+
+    // Heuristic to see if we're invoked as a shebang script, namely,
+    // if we have at least one argument, it's the name of an
+    // executable file, and it starts with "#!".
+    Strings savedArgs;
+    auto runEnv = std::regex_search(programName, std::regex("nix$"));
+    if (runEnv && cmdline.size() > 0) {
+        auto script = *cmdline.begin();
+        try {
+            auto lines = tokenizeString<Strings>(readFile(script), "\n");
+            if (std::regex_search(lines.front(), std::regex("^#!"))) {
+                lines.pop_front();
+                for (auto pos = std::next(cmdline.begin()); pos != cmdline.end();pos++)
+                    savedArgs.push_back(*pos);
+                cmdline.clear();
+
+                for (auto line : lines) {
+                    line = chomp(line);
+
+                    std::smatch match;
+                    if (std::regex_match(line, match, std::regex("^#!\\s*nix (.*)$")))
+                        for (const auto & word : shellwords(match[1].str()))
+                            cmdline.push_back(word);
+                    else {
+                        break;
+                    }
+                }
+                cmdline.push_back(script);
+                for (auto pos = savedArgs.begin(); pos != savedArgs.end();pos++)
+                    cmdline.push_back(*pos);
+            }
+        } catch (SysError &) { }
+    }
     for (auto pos = cmdline.begin(); pos != cmdline.end(); ) {
 
         auto arg = *pos;

src/libutil/args.hh

@@ -20,6 +20,10 @@ public:
        wrong. */
     void parseCmdline(const Strings & cmdline);
 
+    /* Parse the command line with argv0, throwing a UsageError if something
+       goes wrong. */
+    void parseCmdline(const string & argv0, const Strings & cmdline);
+
     /* Return a short one-line description of the command. */
     virtual std::string description() { return ""; }
 

src/libutil/util.cc

@@ -13,6 +13,7 @@
 #include <future>
 #include <iostream>
 #include <mutex>
+#include <regex>
 #include <sstream>
 #include <thread>
 
@@ -1361,6 +1362,48 @@ std::string shellEscape(const std::string & s)
 }
 
 
+std::vector<string> shellwords(const string & s)
+{
+    std::regex whitespace("^(\\s+).*");
+    auto begin = s.cbegin();
+    std::vector<string> res;
+    std::string cur;
+    enum state {
+        sBegin,
+        sQuote
+    };
+    state st = sBegin;
+    auto it = begin;
+    for (; it != s.cend(); ++it) {
+        if (st == sBegin) {
+            std::smatch match;
+            if (regex_search(it, s.cend(), match, whitespace)) {
+                cur.append(begin, it);
+                res.push_back(cur);
+                cur.clear();
+                it = match[1].second;
+                begin = it;
+            }
+        }
+        switch (*it) {
+            case '"':
+                cur.append(begin, it);
+                begin = it + 1;
+                st = st == sBegin ? sQuote : sBegin;
+                break;
+            case '\\':
+                /* perl shellwords mostly just treats the next char as part of the string with no special processing */
+                cur.append(begin, it);
+                begin = ++it;
+                break;
+        }
+    }
+    cur.append(begin, it);
+    if (!cur.empty()) res.push_back(cur);
+    return res;
+}
+
+
 void ignoreException()
 {
     try {

src/libutil/util.hh

@@ -470,6 +470,11 @@ std::string toLower(const std::string & s);
 std::string shellEscape(const std::string & s);
 
 
+/* Recreate the effect of the perl shellwords function, breaking up a
+ * string into arguments like a shell word, including escapes */
+std::vector<string> shellwords(const string & s);
+
+
 /* Exception handling in destructors: print an error message, then
    ignore the exception. */
 void ignoreException();

src/nix-build/nix-build.cc

@@ -30,50 +30,6 @@ using namespace std::string_literals;
 
 extern char * * environ __attribute__((weak));
 
-/* Recreate the effect of the perl shellwords function, breaking up a
- * string into arguments like a shell word, including escapes
- */
-std::vector<string> shellwords(const string & s)
-{
-    std::regex whitespace("^(\\s+).*");
-    auto begin = s.cbegin();
-    std::vector<string> res;
-    std::string cur;
-    enum state {
-        sBegin,
-        sQuote
-    };
-    state st = sBegin;
-    auto it = begin;
-    for (; it != s.cend(); ++it) {
-        if (st == sBegin) {
-            std::smatch match;
-            if (regex_search(it, s.cend(), match, whitespace)) {
-                cur.append(begin, it);
-                res.push_back(cur);
-                cur.clear();
-                it = match[1].second;
-                begin = it;
-            }
-        }
-        switch (*it) {
-            case '"':
-                cur.append(begin, it);
-                begin = it + 1;
-                st = st == sBegin ? sQuote : sBegin;
-                break;
-            case '\\':
-                /* perl shellwords mostly just treats the next char as part of the string with no special processing */
-                cur.append(begin, it);
-                begin = ++it;
-                break;
-        }
-    }
-    cur.append(begin, it);
-    if (!cur.empty()) res.push_back(cur);
-    return res;
-}
-
 static void main_nix_build(int argc, char * * argv)
 {
     auto dryRun = false;

src/nix/main.cc

@@ -266,7 +266,7 @@ void mainWrapped(int argc, char * * argv)
     });
 
     try {
-        args.parseCmdline(argvToStrings(argc, argv));
+        args.parseCmdline(programName, argvToStrings(argc, argv));
     } catch (HelpRequested &) {
         std::vector<std::string> subcommand;
         MultiCommand * command = &args;

tests/flakes.sh

@@ -46,6 +46,7 @@ cat > $flake1Dir/flake.nix <<EOF
 
   outputs = inputs: rec {
     packages.$system.foo = import ./simple.nix;
+    packages.$system.fooScript = (import ./shell.nix {}).foo;
     defaultPackage.$system = packages.$system.foo;
 
     # To test "nix flake init".
@@ -54,8 +55,8 @@ cat > $flake1Dir/flake.nix <<EOF
 }
 EOF
 
-cp ./simple.nix ./simple.builder.sh ./config.nix $flake1Dir/
-git -C $flake1Dir add flake.nix simple.nix simple.builder.sh config.nix
+cp ./simple.nix ./shell.nix ./simple.builder.sh ./config.nix $flake1Dir/
+git -C $flake1Dir add flake.nix simple.nix shell.nix simple.builder.sh config.nix
 git -C $flake1Dir commit -m 'Initial'
 
 cat > $flake2Dir/flake.nix <<EOF
@@ -92,7 +93,17 @@ cat > $nonFlakeDir/README.md <<EOF
 FNORD
 EOF
 
-git -C $nonFlakeDir add README.md
+cat > $nonFlakeDir/shebang.sh <<EOF
+#! $(type -P env) nix
+#! nix --offline shell
+#! nix flake1#fooScript
+#! nix --no-write-lock-file --command bash
+set -e
+foo
+EOF
+chmod +x $nonFlakeDir/shebang.sh
+
+git -C $nonFlakeDir add README.md shebang.sh
 git -C $nonFlakeDir commit -m 'Initial'
 
 # Construct a custom registry, additionally test the --registry flag
@@ -772,3 +783,6 @@ EOF
 git -C $flakeFollowsA add flake.nix
 
 nix flake lock $flakeFollowsA 2>&1 | grep 'this is a security violation'
+
+# Test shebang
+[[ $($nonFlakeDir/shebang.sh) = "foo" ]]