simpler and more robust identification of function parameters for AST profiling

src/ID.cc

@@ -100,20 +100,15 @@ ID::ID(const char* arg_name, IDScope arg_scope, bool arg_is_export) {
     name = util::copy_string(arg_name);
     scope = arg_scope;
     is_export = arg_is_export;
-    is_option = false;
-    is_blank = name && extract_var_name(name) == "_";
-    is_const = false;
-    is_enum_const = false;
-    is_type = false;
     offset = 0;
 
-    if ( is_blank )
+    if ( name && extract_var_name(name) == "_" ) {
+        is_blank = true;
         SetType(base_type(TYPE_ANY));
+    }
 
     opt_info = new IDOptInfo(this);
 
-    infer_return_type = false;
-
     SetLocationInfo(&start_location, &end_location);
 }
 

src/ID.h

@@ -81,7 +81,6 @@ public:
     }
 
     bool IsType() const { return is_type; }
-
     void MakeType() { is_type = true; }
 
     void SetVal(ValPtr v);
@@ -160,9 +159,14 @@ protected:
     const char* name;
     IDScope scope;
     bool is_export;
-    bool infer_return_type;
     TypePtr type;
-    bool is_const, is_enum_const, is_type, is_option, is_blank;
+    bool is_capture = false;
+    bool is_const = false;
+    bool is_enum_const = false;
+    bool is_type = false;
+    bool is_option = false;
+    bool is_blank = false;
+    bool infer_return_type = false;
     int offset;
     ValPtr val;
     AttributesPtr attrs;

src/script_opt/ProfileFunc.cc

@@ -24,11 +24,12 @@ p_hash_type p_hash(const Obj* o) {
 
 ProfileFunc::ProfileFunc(const Func* func, const StmtPtr& body, bool _abs_rec_fields) {
     profiled_func = func;
+    profiled_scope = profiled_func->GetScope();
     profiled_body = body.get();
     abs_rec_fields = _abs_rec_fields;
 
-    auto ft = func->GetType()->AsFuncType();
-    auto& fcaps = ft->GetCaptures();
+    profiled_func_t = cast_intrusive<FuncType>(func->GetType());
+    auto& fcaps = profiled_func_t->GetCaptures();
 
     if ( fcaps ) {
         int offset = 0;
@@ -40,7 +41,7 @@ ProfileFunc::ProfileFunc(const Func* func, const StmtPtr& body, bool _abs_rec_fi
         }
     }
 
-    Profile(ft, body);
+    Profile(profiled_func_t.get(), body);
 }
 
 ProfileFunc::ProfileFunc(const Stmt* s, bool _abs_rec_fields) {
@@ -56,6 +57,9 @@ ProfileFunc::ProfileFunc(const Expr* e, bool _abs_rec_fields) {
 
     if ( e->Tag() == EXPR_LAMBDA ) {
         auto func = e->AsLambdaExpr();
+        ASSERT(func->GetType()->Tag() == TYPE_FUNC);
+        profiled_scope = func->GetScope();
+        profiled_func_t = cast_intrusive<FuncType>(func->GetType());
 
         int offset = 0;
 
@@ -75,6 +79,11 @@ ProfileFunc::ProfileFunc(const Expr* e, bool _abs_rec_fields) {
 
 void ProfileFunc::Profile(const FuncType* ft, const StmtPtr& body) {
     num_params = ft->Params()->NumFields();
+
+    auto& ov = profiled_scope->OrderedVars();
+    for ( int i = 0; i < num_params; ++i )
+        params.insert(ov[i].get());
+
     TrackType(ft);
     body->Traverse(this);
 }
@@ -181,28 +190,10 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e) {
             TrackType(id->GetType());
 
             if ( id->IsGlobal() ) {
-                globals.insert(id);
-                all_globals.insert(id);
-
-                const auto& t = id->GetType();
-                if ( t->Tag() == TYPE_FUNC )
-                    if ( t->AsFuncType()->Flavor() == FUNC_FLAVOR_EVENT )
-                        events.insert(id->Name());
-
+                PreID(id);
                 break;
             }
 
-            // This is a tad ugly.  Unfortunately due to the weird way
-            // that Zeek function *declarations* work, there's no reliable
-            // way to get the list of parameters for a function *definition*,
-            // since they can have different names than what's present in the
-            // declaration.  So we identify them directly, by knowing that
-            // they come at the beginning of the frame ... and being careful
-            // to avoid misconfusing a lambda capture with a low frame offset
-            // as a parameter.
-            if ( captures.count(id) == 0 && id->Offset() < num_params )
-                params.insert(id);
-
             locals.insert(id);
 
             break;
@@ -426,11 +417,6 @@ TraversalCode ProfileFunc::PreExpr(const Expr* e) {
             for ( const auto& i : l->OuterIDs() ) {
                 locals.insert(i);
                 TrackID(i);
-
-                // See above re EXPR_NAME regarding the following
-                // logic.
-                if ( captures.count(i) == 0 && i->Offset() < num_params )
-                    params.insert(i);
             }
 
             // In general, we don't want to recurse into the body.

src/script_opt/ProfileFunc.h

@@ -83,6 +83,7 @@ public:
     // Returns the function, body, or expression profiled.  Each can be
     // null depending on the constructor used.
     const Func* ProfiledFunc() const { return profiled_func; }
+    const ScopePtr& ProfiledScope() const { return profiled_scope; }
     const Stmt* ProfiledBody() const { return profiled_body; }
     const Expr* ProfiledExpr() const { return profiled_expr; }
 
@@ -139,6 +140,7 @@ protected:
     TraversalCode PreStmt(const Stmt*) override;
     TraversalCode PreExpr(const Expr*) override;
     TraversalCode PreID(const ID*) override;
+    TraversalCode PreType(const Type*) override;
 
     // Take note of the presence of a given type.
     void TrackType(const Type* t);
@@ -157,6 +159,8 @@ protected:
     // The function, body, or expression profiled.  Can be null
     // depending on which constructor was used.
     const Func* profiled_func = nullptr;
+    ScopePtr profiled_scope;     // null when not in a full function context
+    FuncTypePtr profiled_func_t; // null when not in a full function context
     const Stmt* profiled_body = nullptr;
     const Expr* profiled_expr = nullptr;