diff --git a/packages/unplugin-typegpu/src/babel.ts b/packages/unplugin-typegpu/src/babel.ts
index 3d2b7a8a03..371323aba7 100644
--- a/packages/unplugin-typegpu/src/babel.ts
+++ b/packages/unplugin-typegpu/src/babel.ts
@@ -87,6 +87,14 @@ function assignMetadata(
       t.variableDeclarator(path.node.id, callExpr),
     ]);
     t.inheritLeadingComments(declaration, path.node);
+
+    if (
+      path.parentPath &&
+      (path.parentPath.isExportNamedDeclaration() || path.parentPath.isExportDefaultDeclaration())
+    ) {
+      t.inheritLeadingComments(declaration, path.parentPath.node);
+      path.parentPath.node.leadingComments = null;
+    }
     replacement = declaration;
   }
 
@@ -94,7 +102,17 @@ function assignMetadata(
     // Hoisting the declaration to the top of the scope
     visibility.unshiftContainer('body', replacement as t.Statement);
     this.alreadyTransformed.add(expression);
-    path.remove();
+
+    const id = t.isFunctionDeclaration(path.node) ? path.node.id : undefined;
+    if (id && path.parentPath.isExportNamedDeclaration()) {
+      path.parentPath.replaceWith(
+        t.exportNamedDeclaration(null, [t.exportSpecifier(t.cloneNode(id), t.cloneNode(id))]),
+      );
+    } else if (id && path.parentPath.isExportDefaultDeclaration()) {
+      path.parentPath.replaceWith(t.exportDefaultDeclaration(t.cloneNode(id)));
+    } else {
+      path.remove();
+    }
   } else {
     path.replaceWith(replacement);
   }
diff --git a/packages/unplugin-typegpu/src/core/factory.ts b/packages/unplugin-typegpu/src/core/factory.ts
index 5d953fa2ec..3005352573 100644
--- a/packages/unplugin-typegpu/src/core/factory.ts
+++ b/packages/unplugin-typegpu/src/core/factory.ts
@@ -69,6 +69,14 @@ function assignMetadata(
     if (t.isExportNamedDeclaration(path.parent)) {
       nodeToOverride = path.parent;
       code = `export ${code}`;
+    } else if (t.isExportDefaultDeclaration(path.parent)) {
+      nodeToOverride = path.parent;
+
+      if (t.isFunctionDeclaration(path.node) && path.node.id) {
+        code = `${code}export default ${path.node.id.name};`;
+      } else {
+        code = `export default ${code};`;
+      }
     }
   }
 
diff --git a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts
index 9a32b06b5a..b43703b553 100644
--- a/packages/unplugin-typegpu/test/use-gpu-directive.test.ts
+++ b/packages/unplugin-typegpu/test/use-gpu-directive.test.ts
@@ -1189,7 +1189,7 @@ describe('replaces function statements marked with "use gpu" in place when condi
   });
 });
 
-test('hoists exported marked function statements', async () => {
+describe('hoists exported marked function statements', () => {
   const code = `\
     console.log(add);
     console.log(mul);
@@ -1205,36 +1205,389 @@ test('hoists exported marked function statements', async () => {
       'use gpu';
       return a * b;
     }
-
   `;
 
-  expect(await rollupTransform(code)).toMatchInlineSnapshot(`
-    "/** MUL */
-    const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) {
-          'use gpu';
-          return __tsover_mul(a, b);
-        }), {
+  test('babel', () => {
+    expect(babelTransform(code)).toMatchInlineSnapshot(`
+      "/** MUL */
+      const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function mul(a, b) {
+        'use gpu';
+
+        return __tsover_mul(a, b);
+      }, {
         v: 1,
         name: "mul",
-        ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","*","b"]]]],"externalNames":[]},
-        externals: () => ({}),
-      }) && $.f)({}));
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "*", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      /** ADD */
+      const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) {
+        'use gpu';
 
-    /** ADD */
-    const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) {
-          'use gpu';
-          return __tsover_add(a, b);
-        }), {
+        return __tsover_add(a, b);
+      }, {
         v: 1,
         name: "add",
-        ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
-        externals: () => ({}),
-      }) && $.f)({}));
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "+", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      console.log(add);
+      console.log(mul);
+      export { add };
+      export { mul };"
+    `);
+  });
+
+  test('rollup', async () => {
+    expect(await rollupTransform(code)).toMatchInlineSnapshot(`
+      "/** MUL */
+      const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function mul(a, b) {
+            'use gpu';
+            return __tsover_mul(a, b);
+          }), {
+          v: 1,
+          name: "mul",
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","*","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
 
+      /** ADD */
+      const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) {
+            'use gpu';
+            return __tsover_add(a, b);
+          }), {
+          v: 1,
+          name: "add",
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+
+      console.log(add);
+          console.log(mul);
+
+      export { add, mul };
+      "
+    `);
+  });
+});
+
+describe('hoists default exported marked function statement', () => {
+  const code = `\
     console.log(add);
-        console.log(mul);
 
-    export { add, mul };
-    "
-  `);
+    /** ADD */
+    export default function add(a, b) {
+      'use gpu';
+      return a + b;
+    }
+  `;
+
+  test('babel', () => {
+    expect(babelTransform(code)).toMatchInlineSnapshot(`
+      "/** ADD */
+      const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function add(a, b) {
+        'use gpu';
+
+        return __tsover_add(a, b);
+      }, {
+        v: 1,
+        name: "add",
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "+", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      console.log(add);
+      export default add;"
+    `);
+  });
+
+  test('rollup', async () => {
+    expect(await rollupTransform(code)).toMatchInlineSnapshot(`
+      "/** ADD */
+      const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function add(a, b) {
+            'use gpu';
+            return __tsover_add(a, b);
+          }), {
+          v: 1,
+          name: "add",
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+      console.log(add);
+
+      export { add as default };
+      "
+    `);
+  });
+});
+
+describe('export marked arrow function', () => {
+  const code = `\
+    export const add = (a, b) => {
+      'use gpu';
+      return a + b;
+    };
+
+    const increment = (n) => {
+      'use gpu';
+      return n + 1;
+    };
+
+    export { increment };
+    
+    const mul = (a, b) => {
+      'use gpu';
+      return a * b;
+    };
+
+    export default mul;
+  `;
+
+  test('babel', () => {
+    expect(babelTransform(code)).toMatchInlineSnapshot(`
+      "export const add = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => {
+        'use gpu';
+
+        return __tsover_add(a, b);
+      }, {
+        v: 1,
+        name: "add",
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "+", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      const increment = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = n => {
+        'use gpu';
+
+        return __tsover_add(n, 1);
+      }, {
+        v: 1,
+        name: "increment",
+        ast: {
+          params: [{
+            type: "i",
+            name: "n"
+          }],
+          body: [0, [[10, [1, "n", "+", [5, "1"]]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      export { increment };
+      const mul = /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => {
+        'use gpu';
+
+        return __tsover_mul(a, b);
+      }, {
+        v: 1,
+        name: "mul",
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "*", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});
+      export default mul;"
+    `);
+  });
+
+  test('rollup', async () => {
+    expect(await rollupTransform(code)).toMatchInlineSnapshot(`
+      "const add = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => {
+            'use gpu';
+            return __tsover_add(a, b);
+          }), {
+          v: 1,
+          name: "add",
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+
+          const increment = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((n) => {
+            'use gpu';
+            return __tsover_add(n, 1);
+          }), {
+          v: 1,
+          name: "increment",
+          ast: {"params":[{"type":"i","name":"n"}],"body":[0,[[10,[1,"n","+",[5,"1"]]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+          
+          const mul = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => {
+            'use gpu';
+            return __tsover_mul(a, b);
+          }), {
+          v: 1,
+          name: "mul",
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","*","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+
+      export { add, mul as default, increment };
+      "
+    `);
+  });
+});
+
+describe('anonymous default export marked function statement', () => {
+  const code = `\
+    export default function (a, b) {
+      'use gpu';
+      return a + b;
+    }
+  `;
+
+  test('babel', () => {
+    expect(babelTransform(code)).toMatchInlineSnapshot(`
+      "export default /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = function (a, b) {
+        'use gpu';
+
+        return __tsover_add(a, b);
+      }, {
+        v: 1,
+        name: undefined,
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "+", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});"
+    `);
+  });
+
+  test('rollup', async () => {
+    expect(await rollupTransform(code)).toMatchInlineSnapshot(`
+      "var _virtual_code = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (function (a, b) {
+            'use gpu';
+            return __tsover_add(a, b);
+          }), {
+          v: 1,
+          name: undefined,
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+
+      export { _virtual_code as default };
+      "
+    `);
+  });
+});
+
+describe('anonymous default export marked arrow function', () => {
+  const code = `\
+    export default (a, b) => {
+      'use gpu';
+      return a + b;
+    }
+  `;
+
+  test('babel', () => {
+    expect(babelTransform(code)).toMatchInlineSnapshot(`
+      "export default /*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = (a, b) => {
+        'use gpu';
+
+        return __tsover_add(a, b);
+      }, {
+        v: 1,
+        name: undefined,
+        ast: {
+          params: [{
+            type: "i",
+            name: "a"
+          }, {
+            type: "i",
+            name: "b"
+          }],
+          body: [0, [[10, [1, "a", "+", "b"]]]],
+          externalNames: []
+        },
+        externals: () => {
+          return {};
+        }
+      }) && $.f)({});"
+    `);
+  });
+
+  test('rollup', async () => {
+    expect(await rollupTransform(code)).toMatchInlineSnapshot(`
+      "var _virtual_code = (/*#__PURE__*/($ => (globalThis.__TYPEGPU_META__ ??= new WeakMap()).set($.f = ((a, b) => {
+            'use gpu';
+            return __tsover_add(a, b);
+          }), {
+          v: 1,
+          name: undefined,
+          ast: {"params":[{"type":"i","name":"a"},{"type":"i","name":"b"}],"body":[0,[[10,[1,"a","+","b"]]]],"externalNames":[]},
+          externals: () => ({}),
+        }) && $.f)({}));
+
+      export { _virtual_code as default };
+      "
+    `);
+  });
 });