lower_subgroup_ops: deal with ptrs

Author: Hugobros3

src/shady/passes/lower_subgroup_ops.c

@@ -41,9 +41,8 @@ static bool is_supported_natively(Context* ctx, const Type* element_type) {
 
 static const Node* build_subgroup_first(Context* ctx, BodyBuilder* bb, const Node* src);
 
-static void build_fn_body(Context* ctx, Node* fn, const Node* param, const Type* t) {
+static const Node* generate(Context* ctx, BodyBuilder* bb, const Node* t, const Node* param) {
     IrArena* a = ctx->rewriter.dst_arena;
-    BodyBuilder* bb = begin_body(a);
     const Type* original_t = t;
     t = get_maybe_nominal_type_body(t);
     switch (is_type(t)) {
@@ -56,11 +55,7 @@ static void build_fn_body(Context* ctx, Node* fn, const Node* param, const Type*
                 const Node* e = gen_extract(bb, param, singleton(uint32_literal(a, i)));
                 elements[i] = build_subgroup_first(ctx, bb, e);
             }
-            fn->payload.fun.body = finish_body(bb, fn_ret(a, (Return) {
-                    .fn = fn,
-                    .args = singleton(composite_helper(a, original_t, nodes(a, element_types.count, elements)))
-            }));
-            return;
+            return composite_helper(a, original_t, nodes(a, element_types.count, elements));
         }
         case Type_Int_TAG: {
             if (t->payload.int_type.width == IntTy64) {
@@ -73,17 +68,30 @@ static void build_fn_body(Context* ctx, Node* fn, const Node* param, const Type*
                 hi = convert_int_zero_extend(bb, it, hi);
                 lo = convert_int_zero_extend(bb, it, lo);
                 hi = gen_primop_e(bb, lshift_op, empty(a), mk_nodes(a, hi, int32_literal(a, 32)));
-                const Node* result = gen_primop_e(bb, or_op, empty(a), mk_nodes(a, lo, hi));
-                fn->payload.fun.body = finish_body(bb, fn_ret(a, (Return) {
-                        .fn = fn,
-                        .args = singleton(result)
-                }));
-                return;
+                return gen_primop_e(bb, or_op, empty(a), mk_nodes(a, lo, hi));
             }
             break;
         }
+        case Type_PtrType_TAG: {
+            param = gen_reinterpret_cast(bb, uint64_type(a), param);
+            return gen_reinterpret_cast(bb, t, generate(ctx, bb, uint64_type(a), param));
+        }
         default: break;
     }
+    return NULL;
+}
+
+static void build_fn_body(Context* ctx, Node* fn, const Node* param, const Type* t) {
+    IrArena* a = ctx->rewriter.dst_arena;
+    BodyBuilder* bb = begin_body(a);
+    const Node* result = generate(ctx, bb, t, param);
+    if (result) {
+        fn->payload.fun.body = finish_body(bb, fn_ret(a, (Return) {
+            .fn = fn,
+            .args = singleton(result)
+        }));
+        return;
+    }
 
     log_string(ERROR, "subgroup_first emulation is not supported for ");
     log_node(ERROR, t);