make Let bind variables

Author: Hugobros3

include/shady/grammar.h

@@ -85,7 +85,7 @@ inline static bool is_nominal(const Node* node) {
     NodeTag tag = node->tag;
     if (node->tag == PrimOp_TAG && has_primop_got_side_effects(node->payload.prim_op.op))
         return true;
-    return tag == Function_TAG || tag == BasicBlock_TAG || tag == Constant_TAG || tag == Param_TAG || tag == GlobalVariable_TAG || tag == NominalType_TAG || tag == Case_TAG;
+    return tag == Function_TAG || tag == BasicBlock_TAG || tag == Constant_TAG || tag == Param_TAG || tag == Variablez_TAG || tag == GlobalVariable_TAG || tag == NominalType_TAG || tag == Case_TAG;
 }
 
 inline static bool is_function(const Node* node) { return node->tag == Function_TAG; }

include/shady/grammar.json

@@ -88,6 +88,10 @@
       "name": "param",
       "generate-enum": false
     },
+    {
+      "name": "variable",
+      "generate-enum": false
+    },
     {
       "name": "instruction"
     },
@@ -254,7 +258,7 @@
     {
       "name": "Variablez",
       "snake_name": "varz",
-      "class": ["value"],
+      "class": ["value", "variable"],
       "constructor": "custom",
       "ops": [
         { "name": "name", "class": "string" },
@@ -448,6 +452,7 @@
       "constructor": "custom",
       "ops": [
         { "name": "instruction", "class": "instruction" },
+        { "name": "variables", "class": "variable", "list": true },
         { "name": "tail", "class": "case" }
       ]
     },
@@ -458,7 +463,8 @@
       "front-end-only": true,
       "ops": [
         { "name": "instruction", "class": "instruction" },
-        { "name": "variables", "class": "param", "list": true }
+        { "name": "variables", "class": "param", "list": true },
+        { "name": "types", "class": "type", "list": true }
       ]
     },
     {
@@ -655,8 +661,7 @@
       "constructor": "custom",
       "class": "case",
       "description": [
-        "An unnamed abstraction that lives inside a function, and can be used as part of various control-flow constructs",
-        "Most notably, the tails of standard `let` nodes"
+        "An unnamed abstraction that lives inside a function, and can be used as part of various control-flow constructs"
       ],
       "ops": [
         { "name": "params", "class": "param", "list": true },

include/shady/ir.h

@@ -225,8 +225,9 @@ const Node* quote_helper(IrArena*, Nodes values);
 const Node* prim_op_helper(IrArena*, Op, Nodes, Nodes);
 
 // terminators
-const Node* let(IrArena*, const Node* instruction, const Node* tail);
-const Node* let_mut(IrArena*, const Node* instruction, Nodes variables);
+const Node* var(IrArena* arena, const char* name, const Node* instruction, size_t i);
+const Node* let(IrArena*, const Node* instruction, Nodes vars, const Node* tail);
+const Node* let_mut(IrArena*, const Node* instruction, Nodes variables, Nodes types);
 const Node* jump_helper(IrArena* a, const Node* dst, Nodes args);
 
 // decl ctors

src/frontends/llvm/l2s.c

@@ -78,7 +78,7 @@ static TodoBB prepare_bb(Parser* p, FnParseCtx* fn_ctx, LLVMBasicBlockRef bb) {
     while (instr) {
         switch (LLVMGetInstructionOpcode(instr)) {
             case LLVMPHI: {
-                const Node* nparam = var(a, qualified_type_helper(convert_type(p, LLVMTypeOf(instr)), false), "phi");
+                const Node* nparam = param(a, qualified_type_helper(convert_type(p, LLVMTypeOf(instr)), false), "phi");
                 insert_dict(LLVMValueRef, const Node*, p->map, instr, nparam);
                 append_list(LLVMValueRef, phis, instr);
                 params = append_nodes(a, params, nparam);
@@ -137,9 +137,9 @@ const Node* convert_function(Parser* p, LLVMValueRef fn) {
     for (LLVMValueRef oparam = LLVMGetFirstParam(fn); oparam; oparam = LLVMGetNextParam(oparam)) {
         LLVMTypeRef ot = LLVMTypeOf(oparam);
         const Type* t = convert_type(p, ot);
-        const Node* param = var(a, qualified_type_helper(t, false), LLVMGetValueName(oparam));
-        insert_dict(LLVMValueRef, const Node*, p->map, oparam, param);
-        params = append_nodes(a, params, param);
+        const Node* nparam = param(a, qualified_type_helper(t, false), LLVMGetValueName(oparam));
+        insert_dict(LLVMValueRef, const Node*, p->map, oparam, nparam);
+        params = append_nodes(a, params, nparam);
         if (oparam == LLVMGetLastParam(fn))
             break;
     }

src/frontends/llvm/l2s_postprocess.c

@@ -39,14 +39,12 @@ static const Node* wrap_in_controls(Context* ctx, Controls* controls, const Node
         const Node* token = controls->tokens.nodes[i];
         const Node* dst = controls->destinations.nodes[i];
         Nodes o_dst_params = get_abstraction_params(dst);
-        LARRAY(const Node*, new_control_params, o_dst_params.count);
-        for (size_t j = 0; j < o_dst_params.count; j++)
-            new_control_params[j] = param(a, o_dst_params.nodes[j]->payload.param.type, unique_name(a, "v"));
-        Nodes nparams = nodes(a, o_dst_params.count, new_control_params);
-        body = let(a, control(a, (Control) {
-            .yield_types = get_param_types(a, o_dst_params),
-            .inside = case_(a, singleton(token), body)
-        }), case_(a, nparams, jump_helper(a, rewrite_node(&ctx->rewriter, dst), nparams)));
+        BodyBuilder* bb = begin_body(a);
+        Nodes results = bind_instruction(bb, control(a, (Control) {
+                .yield_types = get_param_types(a, o_dst_params),
+                .inside = case_(a, singleton(token), body)
+        }));
+        body = finish_body(bb, jump_helper(a, rewrite_node(&ctx->rewriter, dst), results));
     }
     return body;
 }

src/frontends/slim/parser.c

@@ -368,7 +368,7 @@ static void expect_parameters(ctxparams, Nodes* parameters, Nodes* default_value
             const char* id = accept_identifier(ctx);
             expect(id);
 
-            const Node* node = var(arena, qtype, id);
+            const Node* node = param(arena, qtype, id);
             append_list(Node*, params, node);
 
             if (default_values) {
@@ -617,11 +617,11 @@ static const Node* accept_control_flow_instruction(ctxparams, Node* fn) {
             expect(accept_token(ctx, lpar_tok));
             String str = accept_identifier(ctx);
             expect(str);
-            const Node* param = var(arena, join_point_type(arena, (JoinPointType) {
+            const Node* jp = param(arena, join_point_type(arena, (JoinPointType) {
                 .yield_types = yield_types,
             }), str);
             expect(accept_token(ctx, rpar_tok));
-            const Node* body = case_(arena, singleton(param), expect_body(ctx, fn, NULL));
+            const Node* body = case_(arena, singleton(jp), expect_body(ctx, fn, NULL));
             return control(arena, (Control) {
                 .inside = body,
                 .yield_types = yield_types
@@ -726,6 +726,15 @@ static const Node* expect_jump(ctxparams) {
     });
 }
 
+/// for convenience, parse variables as parameters
+static Nodes params2vars(IrArena* arena, const Node* instruction, Nodes params) {
+    LARRAY(const Node*, vars, params.count);
+    for (size_t i = 0; i < params.count; i++) {
+        vars[i] = var(arena, params.nodes[i]->payload.param.name, instruction, i);
+    }
+    return nodes(arena, params.count, vars);
+}
+
 static const Node* accept_terminator(ctxparams, Node* fn) {
     TokenTag tag = curr_token(tokenizer).tag;
     switch (tag) {
@@ -738,7 +747,7 @@ static const Node* accept_terminator(ctxparams, Node* fn) {
                 case let_tok: {
                     const Node* lam = accept_case(ctx, fn);
                     expect(lam);
-                    return let(arena, instruction, lam);
+                    return let(arena, instruction, params2vars(arena, instruction, get_abstraction_params(lam)), get_abstraction_body(lam));
                 }
                 default: SHADY_UNREACHABLE;
             }

src/frontends/spirv/s2s.c

@@ -864,7 +864,7 @@ size_t parse_spv_instruction_at(SpvParser* parser, size_t instruction_offset) {
             parser->defs[result].type = Value;
             String param_name = get_name(parser, result);
             param_name = param_name ? param_name : format_string_arena(parser->arena->arena, "param%d", parser->fun_arg_i);
-            parser->defs[result].node = var(parser->arena, qualified_type_helper(get_def_type(parser, result_t), parser->is_entry_pt), param_name);
+            parser->defs[result].node = param(parser->arena, qualified_type_helper(get_def_type(parser, result_t), parser->is_entry_pt), param_name);
             break;
         }
         case SpvOpLabel: {
@@ -923,7 +923,7 @@ size_t parse_spv_instruction_at(SpvParser* parser, size_t instruction_offset) {
             parser->defs[result].type = Value;
             String phi_name = get_name(parser, result);
             phi_name = phi_name ? phi_name : unique_name(parser->arena, "phi");
-            parser->defs[result].node = var(parser->arena, qualified_type_helper(get_def_type(parser, result_t), false), phi_name);
+            parser->defs[result].node = param(parser->arena, qualified_type_helper(get_def_type(parser, result_t), false), phi_name);
             assert(size % 2 == 1);
             int num_callsites = (size - 3) / 2;
             for (size_t i = 0; i < num_callsites; i++) {

src/shady/analysis/free_variables.c

@@ -23,24 +23,28 @@ typedef struct {
 static void search_op_for_free_variables(Context* visitor, NodeClass class, String op_name, const Node* node) {
     assert(node);
     switch (node->tag) {
-        case Param_TAG: {
-            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;
+        case Let_TAG: {
+            Nodes variables = node->payload.let.variables;
+            for (size_t j = 0; j < variables.count; j++) {
+                const Node* var = variables.nodes[j];
+                bool r = insert_set_get_result(const Node*, visitor->current_scope->bound_set, var);
+                assert(r);
             }
-            insert_set_get_result(const Node*, visitor->current_scope->free_set, node);
             break;
         }
-        case Variablez_TAG: {
-            insert_set_get_result(const Node*, visitor->current_scope->free_set, node);
-            break;
+        case Variablez_TAG:
+        case Param_TAG: {
+            const Node** found = find_key_dict(const Node*, visitor->current_scope->bound_set, node);
+            if (!found)
+                insert_set_get_result(const Node*, visitor->current_scope->free_set, node);
+            return;
         }
         case Function_TAG:
         case Case_TAG:
         case BasicBlock_TAG: assert(false);
-        default: visit_node_operands(&visitor->visitor, IGNORE_ABSTRACTIONS_MASK, node); break;
+        default: break;
     }
+    visit_node_operands(&visitor->visitor, IGNORE_ABSTRACTIONS_MASK | NcVariable, node);
 }
 
 static CFNodeVariables* create_node_variables(CFNode* cfnode) {
@@ -57,11 +61,7 @@ static CFNodeVariables* visit_domtree(Context* ctx, CFNode* cfnode, int depth, C
     ctx = &new_context;
 
     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);
-    }
+    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;
 
@@ -85,15 +85,21 @@ static CFNodeVariables* visit_domtree(Context* ctx, CFNode* cfnode, int depth, C
         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);
+            const Node** found = find_key_dict(const Node*, ctx->current_scope->bound_set, free_var);
+            if (!found)
+                insert_set_get_result(const Node*, ctx->current_scope->free_set, free_var);
             next:;
         }
     }
 
+    if (parent) {
+        size_t j = 0;
+        const Node* bound;
+        while (dict_iter(parent->bound_set, &j, &bound, NULL)) {
+            insert_set_get_result(const Node*, ctx->current_scope->bound_set, bound);
+        }
+    }
+
     /*String abs_name = get_abstraction_name_unsafe(abs);
     for (int i = 0; i < depth; i++)
         debugvv_print(" ");

src/shady/body_builder.c

@@ -23,7 +23,7 @@ BodyBuilder* begin_body(IrArena* a) {
     return bb;
 }
 
-Node* var(IrArena* arena, const char* name, const Node* instruction, size_t i);
+const Node* var(IrArena* arena, const char* name, const Node* instruction, size_t i);
 
 static Nodes create_output_variables(IrArena* a, const Node* value, size_t outputs_count, const Node** output_types, String const output_names[]) {
     Nodes types;
@@ -93,7 +93,7 @@ Nodes bind_instruction_explicit_result_types(BodyBuilder* bb, const Node* instru
 
 Nodes create_mutable_variables(BodyBuilder* bb, const Node* instruction, Nodes provided_types, String const output_names[]) {
     Nodes mutable_vars = create_output_variables(bb->arena, instruction, provided_types.count, provided_types.nodes, output_names);
-    const Node* let_mut_instr = let_mut(bb->arena, instruction, mutable_vars);
+    const Node* let_mut_instr = let_mut(bb->arena, instruction, mutable_vars, provided_types);
     return bind_internal(bb, let_mut_instr, 0, NULL, NULL);
 }
 
@@ -109,12 +109,19 @@ void bind_variables(BodyBuilder* bb, Nodes vars, Nodes values) {
     append_list(StackEntry, bb->stack, entry);
 }
 
+void bind_variables2(BodyBuilder* bb, Nodes vars, const Node* instr) {
+    StackEntry entry = {
+        .instr = instr,
+        .vars = vars,
+    };
+    append_list(StackEntry, bb->stack, entry);
+}
+
 const Node* finish_body(BodyBuilder* bb, const Node* terminator) {
     size_t stack_size = entries_count_list(bb->stack);
     for (size_t i = stack_size - 1; i < stack_size; i--) {
         StackEntry entry = read_list(StackEntry, bb->stack)[i];
-        const Node* lam = case_(bb->arena, entry.vars, terminator);
-        terminator = let(bb->arena, entry.instr, lam);
+        terminator = let(bb->arena, entry.instr, entry.vars, case_(bb->arena, empty(bb->arena), terminator));
     }
 
     destroy_list(bb->stack);

src/shady/compile.c

@@ -67,7 +67,7 @@ CompilationResult run_compiler_passes(CompilerConfig* config, Module** pmod) {
 
     RUN_PASS(specialize_execution_model)
 
-    RUN_PASS(opt_stack)
+    //RUN_PASS(opt_stack)
 
     RUN_PASS(lower_tailcalls)
     RUN_PASS(lower_switch_btree)