Add $# for number of positional arguments

Fixes: #741

docs/reference_guide.md

@@ -52,7 +52,7 @@ This is a work in progress. If something is missing, check the bpftrace source t
     - [6. `nsecs`: Timestamps and Time Deltas](#6-nsecs-timestamps-and-time-deltas)
     - [7. `kstack`: Stack Traces, Kernel](#7-kstack-stack-traces-kernel)
     - [8. `ustack`: Stack Traces, User](#8-ustack-stack-traces-user)
-    - [9. `$1`, ..., `$N`: Positional Parameters](#9-1--n-positional-parameters)
+    - [9. `$1`, ..., `$N`, `$#`: Positional Parameters](#9-1--n--positional-parameters)
 - [Functions](#functions)
     - [1. Builtins](#1-builtins-1)
     - [2. `printf()`: Print Formatted](#2-printf-Printing)
@@ -1131,7 +1131,7 @@ These are special built-in events provided by the bpftrace runtime. `BEGIN` is t
 - `curtask` - Current task struct as a u64
 - `rand` - Random number as a u32
 - `cgroup` - Cgroup ID of the current process
-- `$1`, `$2`, ..., `$N`. - Positional parameters for the bpftrace program
+- `$1`, `$2`, ..., `$N`, `$#`. - Positional parameters for the bpftrace program
 
 Many of these are discussed in other sections (use search).
 
@@ -1378,20 +1378,22 @@ Attaching 1 probe...
 
 Note that for this example to work, bash had to be recompiled with frame pointers.
 
-## 9. `$1`, ..., `$N`: Positional Parameters
+## 9. `$1`, ..., `$N`, `$#`: Positional Parameters
 
-Syntax: `$1`, `$2`, ..., `$N`
+Syntax: `$1`, `$2`, ..., `$N`, `$#`
 
-These are the positional parameters to the bpftrace program, also referred to as command line arguments. If the parameter is numeric (entirerly digits), it can be used as a number. If it is non-numeric, it must be used as a string in the `str()` call. If a parameter is used that was not provided, it will default to zero for numeric context, and "" for string context.
+These are the positional parameters to the bpftrace program, also referred to as command line arguments. If the parameter is numeric (entirely digits), it can be used as a number. If it is non-numeric, it must be used as a string in the `str()` call. If a parameter is used that was not provided, it will default to zero for numeric context, and "" for string context.
+
+`$#` returns the number of positional arguments supplied.
 
 This allows scripts to be written that use basic arguments to change their behavior. If you develop a script that requires more complex argument processing, it may be better suited for bcc instead, which supports Python's argparse and completely custom argument processing.
 
 One-liner example:
 
 ```
-# bpftrace -e 'BEGIN { printf("I got %d, %s\n", $1, str($2)); }' 42 "hello"
+# bpftrace -e 'BEGIN { printf("I got %d, %s (%d args)\n", $1, str($2), $#); }' 42 "hello"
 Attaching 1 probe...
-I got 42, hello
+I got 42, hello (2 args)
 ```
 
 Script example, bsize.d:

src/ast/ast.h

@@ -41,7 +41,8 @@ class Integer : public Expression {
 
 class PositionalParameter : public Expression {
 public:
-  explicit PositionalParameter(long n) : n(n) {}
+  explicit PositionalParameter(PositionalParameterType ptype, long n) : ptype(ptype), n(n) {}
+  PositionalParameterType ptype;
   long n;
   bool is_in_str = false;
 

src/ast/codegen_llvm.cpp

@@ -27,15 +27,28 @@ void CodegenLLVM::visit(Integer &integer)
 
 void CodegenLLVM::visit(PositionalParameter &param)
 {
-  std::string pstr = bpftrace_.get_param(param.n, param.is_in_str);
-  if (bpftrace_.is_numeric(pstr)) {
-    expr_ = b_.getInt64(std::stoll(pstr));
-  } else {
-    Constant *const_str = ConstantDataArray::getString(module_->getContext(), pstr, true);
-    AllocaInst *buf = b_.CreateAllocaBPF(ArrayType::get(b_.getInt8Ty(), pstr.length() + 1), "str");
-    b_.CreateMemSet(buf, b_.getInt8(0), pstr.length() + 1, 1);
-    b_.CreateStore(const_str, buf);
-    expr_ = buf;
+  switch (param.ptype) {
+    case PositionalParameterType::positional:
+      {
+        std::string pstr = bpftrace_.get_param(param.n, param.is_in_str);
+        if (bpftrace_.is_numeric(pstr)) {
+          expr_ = b_.getInt64(std::stoll(pstr));
+        } else {
+          Constant *const_str = ConstantDataArray::getString(module_->getContext(), pstr, true);
+          AllocaInst *buf = b_.CreateAllocaBPF(ArrayType::get(b_.getInt8Ty(), pstr.length() + 1), "str");
+          b_.CreateMemSet(buf, b_.getInt8(0), pstr.length() + 1, 1);
+          b_.CreateStore(const_str, buf);
+          expr_ = buf;
+        }
+      }
+      break;
+    case PositionalParameterType::count:
+      expr_ = b_.getInt64(bpftrace_.num_params());
+      break;
+    default:
+      std::cerr << "unknown parameter type" << std::endl;
+      abort();
+      break;
   }
 }
 

src/ast/printer.cpp

@@ -15,7 +15,17 @@ void Printer::visit(Integer &integer)
 void Printer::visit(PositionalParameter &param)
 {
   std::string indent(depth_, ' ');
-  out_ << indent << "param: $" << param.n << std::endl;
+
+  switch (param.ptype) {
+    case PositionalParameterType::positional:
+      out_ << indent << "param: $" << param.n << std::endl;
+      break;
+    case PositionalParameterType::count:
+      out_ << indent << "param: $#" << std::endl;
+      break;
+    default:
+      break;
+  }
 }
 
 void Printer::visit(String &string)

src/ast/semantic_analyser.cpp

@@ -22,23 +22,35 @@ void SemanticAnalyser::visit(Integer &integer)
 
 void SemanticAnalyser::visit(PositionalParameter &param)
 {
-  if (param.n <= 0) {
-    err_ << "$" << param.n << " is not a valid parameter" << std::endl;
-  }
-
   param.type = SizedType(Type::integer, 8);
-  if (is_final_pass()) {
-    std::string pstr = bpftrace_.get_param(param.n, param.is_in_str);
-    if (!bpftrace_.is_numeric(pstr) && !param.is_in_str) {
-      err_ << "$" << param.n << " used numerically, but given \"" << pstr
-           << "\". Try using str($" << param.n << ")." << std::endl;
-    }
-    if (bpftrace_.is_numeric(pstr) && param.is_in_str) {
-      // This is blocked due to current limitations in our codegen
-      err_ << "$" << param.n << " used in str(), but given numeric value: "
-           << pstr << ". Try $" << param.n << " instead of str($"
-           << param.n << ")." << std::endl;
-    }
+  switch (param.ptype) {
+    case PositionalParameterType::positional:
+      if (param.n <= 0) {
+        err_ << "$" << param.n << " is not a valid parameter" << std::endl;
+      }
+      if (is_final_pass()) {
+        std::string pstr = bpftrace_.get_param(param.n, param.is_in_str);
+        if (!bpftrace_.is_numeric(pstr) && !param.is_in_str) {
+          err_ << "$" << param.n << " used numerically, but given \"" << pstr
+               << "\". Try using str($" << param.n << ")." << std::endl;
+        }
+        if (bpftrace_.is_numeric(pstr) && param.is_in_str) {
+          // This is blocked due to current limitations in our codegen
+          err_ << "$" << param.n << " used in str(), but given numeric value: "
+               << pstr << ". Try $" << param.n << " instead of str($"
+               << param.n << ")." << std::endl;
+        }
+      }
+      break;
+    case PositionalParameterType::count:
+      if (is_final_pass() && param.is_in_str) {
+        err_ << "use $#, not str($#)" << std::endl;
+      }
+      break;
+    default:
+      err_ << "unknown parameter type" << std::endl;
+      param.type = SizedType(Type::none, 0);
+      break;
   }
 }
 

src/bpftrace.cpp

@@ -600,6 +600,11 @@ std::string BPFtrace::get_param(size_t i, bool is_str) const
   return params_.at(i-1);
 }
 
+size_t BPFtrace::num_params() const
+{
+  return params_.size();
+}
+
 void perf_event_lost(void *cb_cookie __attribute__((unused)), uint64_t lost)
 {
   auto bpftrace = static_cast<BPFtrace*>(cb_cookie);

src/bpftrace.h

@@ -93,6 +93,7 @@ class BPFtrace
   void add_param(const std::string &param);
   bool is_numeric(std::string str) const;
   std::string get_param(size_t index, bool is_str) const;
+  size_t num_params() const;
   void request_finalize();
   std::string cmd_;
   int pid_{0};

src/lexer.l

@@ -93,7 +93,6 @@ bpftrace|perf {
 "-"                     { return Parser::make_MINUS(loc); }
 "++"                    { return Parser::make_PLUSPLUS(loc); }
 "--"                    { return Parser::make_MINUSMINUS(loc); }
-"$"                     { return Parser::make_DOLLAR(loc); }
 "*"                     { return Parser::make_MUL(loc); }
 "/"                     { return Parser::make_DIV(loc); }
 "%"                     { return Parser::make_MOD(loc); }
@@ -104,6 +103,8 @@ bpftrace|perf {
 "~"                     { return Parser::make_BNOT(loc); }
 "."                     { return Parser::make_DOT(loc); }
 "->"                    { return Parser::make_PTR(loc); }
+"$"[0-9]+               { return Parser::make_PARAM(yytext, loc); }
+"$"#                    { return Parser::make_PARAMCOUNT(loc); }
 "#"[^!].*               { return Parser::make_CPREPROC(yytext, loc); }
 "if"                    { return Parser::make_IF(yytext, loc); }
 "else"                  { return Parser::make_ELSE(yytext, loc); }

src/parser.yy

@@ -47,6 +47,7 @@ void yyerror(bpftrace::Driver &driver, const char *s);
   QUES       "?"
   ENDPRED    "end predicate"
   COMMA      ","
+  PARAMCOUNT "$#"
   ASSIGN     "="
   EQ         "=="
   NE         "!="
@@ -74,7 +75,6 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 
   MINUS      "-"
   MINUSMINUS "--"
-  DOLLAR     "$"
   MUL        "*"
   DIV        "/"
   MOD        "%"
@@ -96,6 +96,7 @@ void yyerror(bpftrace::Driver &driver, const char *s);
 %token <std::string> STRING "string"
 %token <std::string> MAP "map"
 %token <std::string> VAR "variable"
+%token <std::string> PARAM "positional parameter"
 %token <long> INT "integer"
 %token <std::string> STACK_MODE "stack_mode"
 %nonassoc <std::string> IF "if"
@@ -185,7 +186,9 @@ pred : DIV expr ENDPRED { $$ = new ast::Predicate($2); }
 ternary : expr QUES expr COLON expr { $$ = new ast::Ternary($1, $3, $5); }
      ;
 
-param : DOLLAR INT { $$ = new ast::PositionalParameter($2); }
+param : PARAM      { $$ = new ast::PositionalParameter(PositionalParameterType::positional, std::stoll($1.substr(1, $1.size()-1))); }
+      | PARAMCOUNT { $$ = new ast::PositionalParameter(PositionalParameterType::count, 0); }
+      ;
 
 block : "{" stmts "}"     { $$ = $2; }
       ;

src/types.h

@@ -159,6 +159,12 @@ enum class AsyncAction
 
 uint64_t asyncactionint(AsyncAction a);
 
+enum class PositionalParameterType
+{
+  positional,
+  count
+};
+
 } // namespace bpftrace
 
 

tests/parser.cpp

@@ -52,6 +52,11 @@ TEST(Parser, positional_param)
   test("kprobe:f { $1 }", "Program\n kprobe:f\n  param: $1\n");
 }
 
+TEST(Parser, positional_param_count)
+{
+  test("kprobe:f { $# }", "Program\n kprobe:f\n  param: $#\n");
+}
+
 TEST(Parser, comment)
 {
   test("kprobe:f { /*** ***/0; }", "Program\n kprobe:f\n  int: 0\n");

tests/runtime/other

@@ -68,6 +68,11 @@ RUN bpftrace -e 'BEGIN { printf("-%s-\n", str($1)); exit() }'
 EXPECT --
 TIMEOUT 1
 
+NAME positional arg count
+RUN bpftrace -v -e 'BEGIN { printf("got %d args: %s %d\n", $#, str($1), $2); exit();}' "one" 2
+EXPECT got 2 args: one 2
+TIMEOUT 5
+
 NAME lhist can be cleared
 RUN bpftrace -e 'BEGIN{ @[1] = lhist(3,0,10,1); clear(@); exit() }'
 EXPECT .*

tests/semantic_analyser.cpp

@@ -942,6 +942,9 @@ TEST(semantic_analyser, positional_parameters)
   // Parameters are not required to exist to be used:
   test(bpftrace, "kprobe:f { printf(\"%s\", str($3)); }", 0);
   test(bpftrace, "kprobe:f { printf(\"%d\", $3); }", 0);
+
+  test(bpftrace, "kprobe:f { printf(\"%d\", $#); }", 0);
+  test(bpftrace, "kprobe:f { printf(\"%s\", str($#)); }", 10);
 }
 
 TEST(semantic_analyser, macros)