rewrote free variables analysis to do a whole scope at once

Author: Hugobros3

include/shady/ir.h

@@ -139,6 +139,7 @@ Nodes filter_out_annotation(IrArena*, Nodes, const char* name);
 
 bool        is_abstraction        (const Node*);
 String      get_abstraction_name  (const Node* abs);
+String      get_abstraction_name_unsafe(const Node* abs);
 const Node* get_abstraction_body  (const Node* abs);
 Nodes       get_abstraction_params(const Node* abs);
 

src/common/dict.c

@@ -334,3 +334,11 @@ KeyHash hash_murmur(const void* data, size_t size) {
     final ^= out[3];
     return final;
 }
+
+KeyHash hash_ptr(void** p) {
+    return hash_murmur(p, sizeof(void*));
+}
+
+bool compare_ptrs(void** a, void** b) {
+    return *a == *b;
+}

src/common/dict.h

@@ -46,4 +46,7 @@ bool insert_dict_and_get_result_impl(struct Dict*, void* key, void* value);
 
 KeyHash hash_murmur(const void* data, size_t size);
 
+KeyHash hash_ptr(void**);
+bool compare_ptrs(void**, void**);
+
 #endif

src/shady/analysis/free_variables.c

@@ -8,27 +8,28 @@
 #include "list.h"
 #include "dict.h"
 
+#include <stdlib.h>
 #include <assert.h>
 
 KeyHash hash_node(Node**);
 bool compare_node(Node**, Node**);
 
 typedef struct {
     Visitor visitor;
-    struct Dict* bound_set;
-    struct Dict* set;
-    struct List* free_list;
+    struct Dict* map;
+    CFNodeVariables* current_scope;
 } Context;
 
 static void search_op_for_free_variables(Context* visitor, NodeClass class, String op_name, const Node* node) {
     assert(node);
     switch (node->tag) {
         case Variable_TAG: {
-            if (find_key_dict(const Node*, visitor->bound_set, node))
-                return;
-            if (insert_set_get_result(const Node*, visitor->set, node)) {
-                append_list(const Node*, visitor->free_list, node);
+            Nodes params = get_abstraction_params(visitor->current_scope->node->node);
+            for (size_t i = 0; i < params.count; i++) {
+                if (params.nodes[i] == node)
+                    return;
             }
+            insert_set_get_result(const Node*, visitor->current_scope->free_set, node);
             break;
         }
         case Function_TAG:
@@ -38,26 +39,36 @@ static void search_op_for_free_variables(Context* visitor, NodeClass class, Stri
     }
 }
 
-static void visit_domtree(Context* ctx, CFNode* cfnode, int depth) {
-    const Node* abs = cfnode->node;
+static CFNodeVariables* create_node_variables(CFNode* cfnode) {
+    CFNodeVariables* v = calloc(sizeof(CFNodeVariables), 1);
+    *v = (CFNodeVariables) {
+        .node = cfnode,
+        .free_set = new_set(const Node*, (HashFn) hash_node, (CmpFn) compare_node)
+    };
+    return v;
+}
 
-    bool is_named = abs->tag != Case_TAG;
+static CFNodeVariables* visit_domtree(Context* ctx, CFNode* cfnode, int depth, CFNodeVariables* parent) {
+    Context new_context = *ctx;
+    ctx = &new_context;
 
-    if (is_named) {
-        for (int i = 0; i < depth; i++)
-            debugvv_print(" ");
-        debugvv_print("%s\n", get_abstraction_name(abs));
+    ctx->current_scope = create_node_variables(cfnode);
+    if (parent) {
+        ctx->current_scope->bound_set = clone_dict(parent->bound_set);
+    } else {
+        ctx->current_scope->bound_set = new_set(const Node*, (HashFn) hash_node, (CmpFn) compare_node);
     }
+    insert_dict(CFNode*, CFNodeVariables*, ctx->map, cfnode, ctx->current_scope);
+    const Node* abs = cfnode->node;
+
+    bool is_named = abs->tag != Case_TAG;
 
     // Bind parameters
     Nodes params = get_abstraction_params(abs);
     for (size_t j = 0; j < params.count; j++) {
         const Node* param = params.nodes[j];
-        bool r = insert_set_get_result(const Node*, ctx->bound_set, param);
-        // assert(r);
-        // this can happen if you visit the domtree of a CFG starting _inside_ a loop
-        // we will meet some unbound params but eventually we'll enter their definition after the fact
-        // those params should still be considered free in this case.
+        bool r = insert_set_get_result(const Node*, ctx->current_scope->bound_set, param);
+        assert(r);
     }
 
     const Node* body = get_abstraction_body(abs);
@@ -66,35 +77,83 @@ static void visit_domtree(Context* ctx, CFNode* cfnode, int depth) {
 
     for (size_t i = 0; i < entries_count_list(cfnode->dominates); i++) {
         CFNode* child = read_list(CFNode*, cfnode->dominates)[i];
-        visit_domtree(ctx, child, depth + (is_named ? 1 : 0));
+        CFNodeVariables* child_variables = visit_domtree(ctx, child, depth + (is_named ? 1 : 1), ctx->current_scope);
+        size_t j = 0;
+        const Node* free_var;
+        while (dict_iter(child_variables->free_set, &j, &free_var, NULL)) {
+            for (size_t k = 0; k < params.count; k++) {
+                if (params.nodes[k] == free_var)
+                    goto next;
+            }
+            insert_set_get_result(const Node*, ctx->current_scope->free_set, free_var);
+            next:;
+        }
     }
 
+    /*String abs_name = get_abstraction_name_unsafe(abs);
+    for (int i = 0; i < depth; i++)
+        debugvv_print(" ");
+    if (abs_name)
+        debugvv_print("%s: ", abs_name);
+    else
+        debugvv_print("%%%d: ", abs->id);
+
+    if (true) {
+        bool prev = false;
+        size_t i = 0;
+        const Node* free_var;
+        while (dict_iter(ctx->current_scope->free_set, &i, &free_var, NULL)) {
+            if (prev) {
+                debugvv_print(", ");
+            }
+            log_node(DEBUGVV, free_var);
+            prev = true;
+        }
+    }
+
+    if (true) {
+        debugvv_print(". Binds: ");
+        bool prev = false;
+        for (size_t i = 0; i < params.count; i++) {
+            if (prev) {
+                debugvv_print(", ");
+            }
+            log_node(DEBUGVV, params.nodes[i]);
+            prev = true;
+        }
+    }
+
+    debugvv_print("\n");*/
+
     // Unbind parameters
     for (size_t j = 0; j < params.count; j++) {
         const Node* param = params.nodes[j];
-        bool r = remove_dict(const Node*, ctx->bound_set, param);
+        bool r = remove_dict(const Node*, ctx->current_scope->bound_set, param);
         assert(r);
     }
-}
 
-struct List* compute_free_variables(const Scope* scope, const Node* at) {
-    struct Dict* bound_set = new_set(const Node*, (HashFn) hash_node, (CmpFn) compare_node);
-    struct Dict* set = new_set(const Node*, (HashFn) hash_node, (CmpFn) compare_node);
-    struct List* free_list = new_list(const Node*);
+    return ctx->current_scope;
+}
 
+struct Dict* compute_scope_variables_map(const Scope* scope) {
     Context ctx = {
         .visitor = {
             .visit_op_fn = (VisitOpFn) search_op_for_free_variables,
         },
-        .bound_set = bound_set,
-        .set = set,
-        .free_list = free_list,
+        .map = new_dict(CFNode*, CFNodeVariables*, (HashFn) hash_ptr, (CmpFn) compare_ptrs),
     };
 
-    debugv_print("Computing free variables...\n");
-    visit_domtree(&ctx, scope_lookup(scope, at), 0);
+    debugv_print("Computing free variables for function '%s' ...\n", get_abstraction_name(scope->entry->node));
+    visit_domtree(&ctx, scope->entry, 0, NULL);
+    return ctx.map;
+}
 
-    destroy_dict(bound_set);
-    destroy_dict(set);
-    return free_list;
+void destroy_scope_variables_map(struct Dict* map) {
+    size_t i = 0;
+    CFNodeVariables* value;
+    while (dict_iter(map, &i, NULL, &value)) {
+        destroy_dict(value->bound_set);
+        destroy_dict(value->free_set);
+        free((void*) value);
+    }
 }

src/shady/analysis/free_variables.h

@@ -4,7 +4,15 @@
 #include "shady/ir.h"
 
 typedef struct Scope_ Scope;
+typedef struct CFNode_ CFNode;
 
-struct List* compute_free_variables(const Scope* scope, const Node*);
+typedef struct {
+    CFNode* node;
+    struct Dict* bound_set;
+    struct Dict* free_set;
+} CFNodeVariables;
+
+struct Dict* compute_scope_variables_map(const Scope* scope);
+void destroy_scope_variables_map(struct Dict*);
 
 #endif

src/shady/analysis/scope.c

@@ -550,9 +550,10 @@ static void dump_cf_node(FILE* output, const CFNode* n) {
 
     label = format_string_arena(bb->arena->arena, "%s%s", label, node_tags[body->tag]);
 
-    if (is_basic_block(bb)) {
-        label = format_string_arena(bb->arena->arena, "%s\n%s", get_abstraction_name(bb), label);
-    }
+    String abs_name = get_abstraction_name_unsafe(bb);
+    if (!abs_name)
+        abs_name = format_string_interned(bb->arena, "%%%d", bb->id);
+    label = format_string_arena(bb->arena->arena, "%s\n%s", abs_name, label);
 
     fprintf(output, "bb_%zu [label=\"%s\", color=\"%s\", shape=box];\n", (size_t) n, label, color);
 

src/shady/analysis/verify.c

@@ -46,13 +46,16 @@ static void verify_scoping(Module* mod) {
     struct List* scopes = build_scopes(mod);
     for (size_t i = 0; i < entries_count_list(scopes); i++) {
         Scope* scope = read_list(Scope*, scopes)[i];
-        struct List* leaking = compute_free_variables(scope, scope->entry->node);
-        for (size_t j = 0; j < entries_count_list(leaking); j++) {
-            log_node(ERROR, read_list(const Node*, leaking)[j]);
+        struct Dict* map = compute_scope_variables_map(scope);
+        CFNodeVariables* entry_vars = *find_value_dict(CFNode*, CFNodeVariables*, map, scope->entry);
+        size_t j = 0;
+        const Node* leaking;
+        while (dict_iter(entry_vars->free_set, &j, &leaking, NULL)) {
+            log_node(ERROR, leaking);
             error_print("\n");
         }
-        assert(entries_count_list(leaking) == 0);
-        destroy_list(leaking);
+        assert(entries_count_dict(entry_vars->free_set) == 0);
+        destroy_scope_variables_map(map);
         destroy_scope(scope);
     }
     destroy_list(scopes);

src/shady/node.c

@@ -263,6 +263,16 @@ String get_abstraction_name(const Node* abs) {
     }
 }
 
+String get_abstraction_name_unsafe(const Node* abs) {
+    assert(is_abstraction(abs));
+    switch (abs->tag) {
+        case Function_TAG: return abs->payload.fun.name;
+        case BasicBlock_TAG: return abs->payload.basic_block.name;
+        case Case_TAG: return NULL;
+        default: assert(false);
+    }
+}
+
 const Node* get_abstraction_body(const Node* abs) {
     assert(is_abstraction(abs));
     switch (abs->tag) {

src/shady/passes/lcssa.c

@@ -17,6 +17,7 @@ typedef struct Context_ {
     const Node* current_fn;
     Scope* scope;
     const UsesMap* scope_uses;
+    struct Dict* scope_vars;
     LoopTree* loop_tree;
     struct Dict* lifted_arguments;
 } Context;
@@ -55,9 +56,11 @@ void find_liftable_loop_values(Context* ctx, const Node* old, Nodes* nparams, No
     *lparams = empty(a);
     *nargs = empty(a);
 
-    struct List* fvs = compute_free_variables(ctx->scope, old);
-    for (size_t i = 0; i < entries_count_list(fvs); i++) {
-        const Node* fv = read_list(const Node*, fvs)[i];
+    const Node* fv;
+    size_t i = 0;
+    CFNode* cf_node = scope_lookup(ctx->scope, old);
+    CFNodeVariables* node_vars = *find_value_dict(CFNode*, CFNodeVariables*, ctx->scope_vars, cf_node);
+    while (dict_iter(node_vars->free_set, &i, &fv, NULL)) {
         const Node* defining_abs = get_var_binding_abstraction(ctx->scope_uses, fv);
         const CFNode* defining_cf_node = scope_lookup(ctx->scope, defining_abs);
         assert(defining_cf_node);
@@ -74,7 +77,6 @@ void find_liftable_loop_values(Context* ctx, const Node* old, Nodes* nparams, No
             *nargs = append_nodes(a, *nargs, narg);
         }
     }
-    destroy_list(fvs);
 
     if (nparams->count > 0)
         insert_dict(const Node*, Nodes, ctx->lifted_arguments, old, *nparams);
@@ -152,6 +154,7 @@ const Node* process_node(Context* ctx, const Node* old) {
             ctx->current_fn = old;
             ctx->scope = new_scope(old);
             ctx->scope_uses = create_uses_map(old, (NcDeclaration | NcType));
+            ctx->scope_vars = compute_scope_variables_map(ctx->scope);
             ctx->loop_tree = build_loop_tree(ctx->scope);
 
             Node* new = recreate_decl_header_identity(&ctx->rewriter, old);
@@ -160,6 +163,7 @@ const Node* process_node(Context* ctx, const Node* old) {
             destroy_loop_tree(ctx->loop_tree);
             destroy_uses_map(ctx->scope_uses);
             destroy_scope(ctx->scope);
+            destroy_scope_variables_map(ctx->scope_vars);
             return new;
         }
         case Jump_TAG: {

src/shady/passes/lift_indirect_targets.c

@@ -26,6 +26,7 @@ typedef struct Context_ {
     Rewriter rewriter;
     Scope* scope;
     const UsesMap* scope_uses;
+    struct Dict* scope_vars;
 
     struct Dict* lifted;
     bool disable_lowering;
@@ -78,19 +79,23 @@ static LiftedCont* lambda_lift(Context* ctx, const Node* cont, String given_name
     String name = is_basic_block(cont) ? format_string_arena(a->arena, "%s_%s", get_abstraction_name(cont->payload.basic_block.fn), get_abstraction_name(cont)) : unique_name(a, given_name);
 
     // Compute the live stuff we'll need
-    Scope* scope = new_scope(cont);
-    struct List* recover_context = compute_free_variables(scope, cont);
-    size_t recover_context_size = entries_count_list(recover_context);
-    destroy_scope(scope);
-
-    debugv_print("free (spilled) variables at '%s': ", name);
-    for (size_t i = 0; i < recover_context_size; i++) {
-        const Node* item = read_list(const Node*, recover_context)[i];
-        debugv_print(get_value_name_safe(item));
-        if (i + 1 < recover_context_size)
-            debugv_print(", ");
+    CFNode* cf_node = scope_lookup(ctx->scope, cont);
+    CFNodeVariables* node_vars = *find_value_dict(CFNode*, CFNodeVariables*, ctx->scope_vars, cf_node);
+    struct List* recover_context = new_list(const Node*);
+    size_t recover_context_size = entries_count_dict(node_vars->free_set);
+
+    {
+        debugv_print("lambda_lift: free (to-be-spilled) variables at '%s' (count=%d): ", name, entries_count_dict(node_vars->free_set));
+        size_t i = 0;
+        const Node* item;
+        while (dict_iter(node_vars->free_set, &i, &item, NULL)) {
+            append_list(const Node*, recover_context, item );
+            debugv_print(get_value_name_safe(item));
+            if (i + 1 < recover_context_size)
+                debugv_print(", ");
+        }
+        debugv_print("\n");
     }
-    debugv_print("\n");
 
     // Create and register new parameters for the lifted continuation
     Nodes new_params = recreate_variables(&ctx->rewriter, oparams);
@@ -166,12 +171,14 @@ static const Node* process_node(Context* ctx, const Node* node) {
             Context fn_ctx = *ctx;
             fn_ctx.scope = new_scope(node);
             fn_ctx.scope_uses = create_uses_map(node, (NcDeclaration | NcType));
+            fn_ctx.scope_vars = compute_scope_variables_map(fn_ctx.scope);
             ctx = &fn_ctx;
 
             Node* new = recreate_decl_header_identity(&ctx->rewriter, node);
             recreate_decl_body_identity(&ctx->rewriter, node, new);
 
             destroy_uses_map(ctx->scope_uses);
+            destroy_scope_variables_map(ctx->scope_vars);
             destroy_scope(ctx->scope);
             return new;
         }