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"
    }
    "$@"
```

fix: include programName usage

EDIT: For posterity I've changed shellwords to shellwords2 in order
      not to interfere with other changes during a rebase.
      shellwords2 is removed in a later commit. -- roberth

src/libutil/args.cc

@@ -6,6 +6,7 @@
 #include "users.hh"
 #include "json-utils.hh"
 
+#include <regex>
 #include <glob.h>
 
 namespace nix {
@@ -78,6 +79,12 @@ std::optional<std::string> RootArgs::needsCompletion(std::string_view s)
 }
 
 void RootArgs::parseCmdline(const Strings & _cmdline)
+{
+    // Default via 5.1.2.2.1 in C standard
+    Args::parseCmdline("", _cmdline);
+}
+
+void Args::parseCmdline(const std::string & programName, const Strings & _cmdline)
 {
     Strings pendingArgs;
     bool dashDash = false;
@@ -93,6 +100,36 @@ void RootArgs::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 isNixCommand = std::regex_search(programName, std::regex("nix$"));
+    if (isNixCommand && 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\\s(.*)$")))
+                        for (const auto & word : shellwords(match[1].str()))
+                            cmdline.push_back(word);
+                }
+                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

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

src/libutil/util.cc

@@ -5,6 +5,8 @@
 #include <cctype>
 #include <iostream>
 #include <grp.h>
+#include <regex>
+
 
 namespace nix {
 
@@ -136,6 +138,49 @@ std::string shellEscape(const std::string_view s)
     return r;
 }
 
+/* Recreate the effect of the perl shellwords function, breaking up a
+ * string into arguments like a shell word, including escapes
+ */
+std::vector<std::string> shellwords2(const std::string & s)
+{
+    std::regex whitespace("^(\\s+).*");
+    auto begin = s.cbegin();
+    std::vector<std::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(Verbosity lvl)
 {

src/libutil/util.hh

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

src/nix/main.cc

@@ -428,7 +428,7 @@ void mainWrapped(int argc, char * * argv)
     });
 
     try {
-        args.parseCmdline(argvToStrings(argc, argv));
+        args.parseCmdline(programName, argvToStrings(argc, argv));
     } catch (UsageError &) {
         if (!args.helpRequested && !args.completions) throw;
     }

tests/functional/flakes/common.sh

@@ -11,6 +11,7 @@ writeSimpleFlake() {
   outputs = inputs: rec {
     packages.$system = rec {
       foo = import ./simple.nix;
+      fooScript = (import ./shell.nix {}).foo;
       default = foo;
     };
     packages.someOtherSystem = rec {
@@ -24,13 +25,13 @@ writeSimpleFlake() {
 }
 EOF
 
-    cp ../simple.nix ../simple.builder.sh ../config.nix $flakeDir/
+    cp ../simple.nix ../shell.nix ../simple.builder.sh ../config.nix $flakeDir/
 }
 
 createSimpleGitFlake() {
     local flakeDir="$1"
     writeSimpleFlake $flakeDir
-    git -C $flakeDir add flake.nix simple.nix simple.builder.sh config.nix
+    git -C $flakeDir add flake.nix simple.nix shell.nix simple.builder.sh config.nix
     git -C $flakeDir commit -m 'Initial'
 }
 

tests/functional/flakes/flakes.sh

@@ -66,7 +66,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
@@ -511,3 +521,6 @@ nix flake metadata "$flake2Dir" --reference-lock-file $TEST_ROOT/flake2-overridd
 
 # reference-lock-file can only be used if allow-dirty is set.
 expectStderr 1 nix flake metadata "$flake2Dir" --no-allow-dirty --reference-lock-file $TEST_ROOT/flake2-overridden.lock
+
+# Test shebang
+[[ $($nonFlakeDir/shebang.sh) = "foo" ]]