feat: std.copy (#2489)

Author: aleksanderkatan

apps/typegpu-docs/src/content/docs/apis/data-schemas.mdx

@@ -514,7 +514,7 @@ WGSL needs the `unrestricted_pointer_parameters` extension.
 This extension is enabled automatically when available.
 :::
 
-## Copy constructors
+## References and copying data
 
 One of the biggest differences between JavaScript and WGSL is the existence of value/reference types.
 Let's take a look at a simple TypeGPU function and the WGSL code it resolves to.
@@ -545,36 +545,32 @@ fn myFn(){
 }
 ```
 
-In JavaScript, `s2` is a reference to `s1` and therefore `s1` is modified when `s2` is modified.
-On the WGSL side, we make `s2` a reference to `s1` and modifying `s2.n` is translated to modifying `(*s2).n`.
+Even though modifying `s2` results in a change to `s1` in both cases, it happens in two entirely different ways:
 
-Schemas can be used to create deep copies of values.
+- In JavaScript, `s2` is a reference to `s1`, so modifying `s2` also modifies `s1`.
+- On the WGSL side, we make `s2` a pointer to `s1`, so modifying `s2.n` is translated to modifying `(*s2).n`.
 
+To force an explicit copy in both cases, use a schema when the exact type is known:
 ```ts twoslash
-import tgpu, { d } from 'typegpu';
+import tgpu, { d, std } from 'typegpu';
 // ---cut---
 const MyStruct = d.struct({ n: d.u32 });
 
 function myFn() {
   'use gpu';
-  const s1 = MyStruct({ n: 0 });
-  const s2 = MyStruct(s1); // deep copy of s1
-  s2.n = 1; // s1 is not modified
+  const struct = MyStruct({ n: 0 })
+  const copy = MyStruct(struct);
 }
 ```
 
-The TypeGPU shader generator understands this and generates a simple assignment, which copies the value of `s1` into `s2`.
-
-```wgsl
-// WGSL
-struct MyStruct {
-  n: u32,
-}
-
-fn myFn(){
-  var s1 = MyStruct(0);
-  var s2 = s1; // copies s1 into s2
-  s2.n = 1; // s1 is not modified
+Or use `std.copy` when the schema may vary:
+```ts twoslash
+import tgpu, { d, std } from 'typegpu';
+// ---cut---
+function myFn(arg: d.v2u | d.v3u | d.v4u) {
+  'use gpu';
+  const vec = std.copy(arg);
+  vec.x++; 
 }
 ```
 

packages/typegpu/src/std/copy.ts

@@ -0,0 +1,34 @@
+import { dualImpl } from '../core/function/dualImpl.ts';
+import { stitch } from '../core/resolve/stitch.ts';
+import { isMatInstance, isVecInstance, WORKAROUND_getSchema } from '../data/wgslTypes.ts';
+
+function cpuCopy<T>(e: T): T {
+  if (isVecInstance(e) || isMatInstance(e)) {
+    const schema = WORKAROUND_getSchema(e);
+    return schema(e as never) as T;
+  }
+
+  if (Array.isArray(e)) {
+    return e.map(cpuCopy) as T;
+  }
+
+  if (typeof e === 'object' && e !== null) {
+    return Object.fromEntries(Object.entries(e).map(([key, value]) => [key, cpuCopy(value)])) as T;
+  }
+
+  return e;
+}
+
+export const copy = dualImpl({
+  name: 'copy',
+  signature: (arg) => {
+    return {
+      argTypes: [arg],
+      returnType: arg,
+    };
+  },
+  normalImpl: cpuCopy,
+  codegenImpl(_ctx, [a]) {
+    return stitch`${a}`;
+  },
+});

packages/typegpu/src/std/index.ts

@@ -4,6 +4,8 @@
 
 // NOTE: This is a barrel file, internal files should not import things from this file
 
+export { copy } from './copy.ts';
+
 export { discard } from './discard.ts';
 
 export {