mirror of
https://github.com/storytold/spark.git
synced 2026-10-18 07:40:38 +00:00
Initial commit
This commit is contained in:
Diego Marcos
@@ -0,0 +1,593 @@
|
||||
import { Dyno, UnaryOp, unindent, unindentLines } from "./base";
|
||||
import {
|
||||
DynoOutput,
|
||||
type DynoVal,
|
||||
type DynoValue,
|
||||
type HasDynoOut,
|
||||
} from "./value";
|
||||
|
||||
export const Gsplat = { type: "Gsplat" } as { type: "Gsplat" };
|
||||
export const TPackedSplats = { type: "PackedSplats" } as {
|
||||
type: "PackedSplats";
|
||||
};
|
||||
|
||||
export const numPackedSplats = (
|
||||
packedSplats: DynoVal<typeof TPackedSplats>,
|
||||
): DynoVal<"int"> => new NumPackedSplats({ packedSplats });
|
||||
export const readPackedSplat = (
|
||||
packedSplats: DynoVal<typeof TPackedSplats>,
|
||||
index: DynoVal<"int">,
|
||||
): DynoVal<typeof Gsplat> => new ReadPackedSplat({ packedSplats, index });
|
||||
export const readPackedSplatRange = (
|
||||
packedSplats: DynoVal<typeof TPackedSplats>,
|
||||
index: DynoVal<"int">,
|
||||
base: DynoVal<"int">,
|
||||
count: DynoVal<"int">,
|
||||
): DynoVal<typeof Gsplat> =>
|
||||
new ReadPackedSplatRange({ packedSplats, index, base, count });
|
||||
export const splitGsplat = (gsplat: DynoVal<typeof Gsplat>) =>
|
||||
new SplitGsplat({ gsplat });
|
||||
export const combineGsplat = ({
|
||||
gsplat,
|
||||
flags,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
}: {
|
||||
gsplat?: DynoVal<typeof Gsplat>;
|
||||
flags?: DynoVal<"uint">;
|
||||
index?: DynoVal<"int">;
|
||||
center?: DynoVal<"vec3">;
|
||||
scales?: DynoVal<"vec3">;
|
||||
quaternion?: DynoVal<"vec4">;
|
||||
rgba?: DynoVal<"vec4">;
|
||||
rgb?: DynoVal<"vec3">;
|
||||
opacity?: DynoVal<"float">;
|
||||
x?: DynoVal<"float">;
|
||||
y?: DynoVal<"float">;
|
||||
z?: DynoVal<"float">;
|
||||
r?: DynoVal<"float">;
|
||||
g?: DynoVal<"float">;
|
||||
b?: DynoVal<"float">;
|
||||
}): DynoVal<typeof Gsplat> => {
|
||||
return new CombineGsplat({
|
||||
gsplat,
|
||||
flags,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
});
|
||||
};
|
||||
export const gsplatNormal = (gsplat: DynoVal<typeof Gsplat>): DynoVal<"vec3"> =>
|
||||
new GsplatNormal({ gsplat });
|
||||
|
||||
export const transformGsplat = (
|
||||
gsplat: DynoVal<typeof Gsplat>,
|
||||
{
|
||||
scale,
|
||||
rotate,
|
||||
translate,
|
||||
recolor,
|
||||
}: {
|
||||
scale?: DynoVal<"float">;
|
||||
rotate?: DynoVal<"vec4">;
|
||||
translate?: DynoVal<"vec3">;
|
||||
recolor?: DynoVal<"vec4">;
|
||||
},
|
||||
): DynoVal<typeof Gsplat> => {
|
||||
return new TransformGsplat({ gsplat, scale, rotate, translate, recolor });
|
||||
};
|
||||
|
||||
export const defineGsplat = unindent(`
|
||||
struct Gsplat {
|
||||
vec3 center;
|
||||
uint flags;
|
||||
vec3 scales;
|
||||
int index;
|
||||
vec4 quaternion;
|
||||
vec4 rgba;
|
||||
};
|
||||
const uint GSPLAT_FLAG_ACTIVE = 1u << 0u;
|
||||
|
||||
bool isGsplatActive(uint flags) {
|
||||
return (flags & GSPLAT_FLAG_ACTIVE) != 0u;
|
||||
}
|
||||
`);
|
||||
|
||||
export const definePackedSplats = unindent(`
|
||||
struct PackedSplats {
|
||||
usampler2DArray texture;
|
||||
int numSplats;
|
||||
};
|
||||
`);
|
||||
|
||||
export class NumPackedSplats extends UnaryOp<
|
||||
typeof TPackedSplats,
|
||||
"int",
|
||||
"numSplats"
|
||||
> {
|
||||
constructor({
|
||||
packedSplats,
|
||||
}: { packedSplats: DynoVal<typeof TPackedSplats> }) {
|
||||
super({ a: packedSplats, outKey: "numSplats", outTypeFunc: () => "int" });
|
||||
this.statements = ({ inputs, outputs }) => [
|
||||
`${outputs.numSplats} = ${inputs.a}.numSplats;`,
|
||||
];
|
||||
}
|
||||
}
|
||||
|
||||
const defineReadPackedSplat = unindent(`
|
||||
bool readPackedSplat(usampler2DArray texture, int numSplats, int index, out Gsplat gsplat) {
|
||||
if ((index >= 0) && (index < numSplats)) {
|
||||
uvec4 packed = texelFetch(texture, splatTexCoord(index), 0);
|
||||
unpackSplat(packed, gsplat.center, gsplat.scales, gsplat.quaternion, gsplat.rgba);
|
||||
return true;
|
||||
} else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
`);
|
||||
|
||||
export class ReadPackedSplat
|
||||
extends Dyno<
|
||||
{ packedSplats: typeof TPackedSplats; index: "int" },
|
||||
{ gsplat: typeof Gsplat }
|
||||
>
|
||||
implements HasDynoOut<typeof Gsplat>
|
||||
{
|
||||
constructor({
|
||||
packedSplats,
|
||||
index,
|
||||
}: { packedSplats?: DynoVal<typeof TPackedSplats>; index?: DynoVal<"int"> }) {
|
||||
super({
|
||||
inTypes: { packedSplats: TPackedSplats, index: "int" },
|
||||
outTypes: { gsplat: Gsplat },
|
||||
inputs: { packedSplats, index },
|
||||
globals: () => [defineGsplat, definePackedSplats, defineReadPackedSplat],
|
||||
statements: ({ inputs, outputs }) => {
|
||||
const { gsplat } = outputs;
|
||||
if (!gsplat) {
|
||||
return [];
|
||||
}
|
||||
const { packedSplats, index } = inputs;
|
||||
let statements: string[];
|
||||
if (packedSplats && index) {
|
||||
statements = unindentLines(`
|
||||
if (readPackedSplat(${packedSplats}.texture, ${packedSplats}.numSplats, ${index}, ${gsplat})) {
|
||||
bool zeroSize = all(equal(${gsplat}.scales, vec3(0.0, 0.0, 0.0)));
|
||||
${gsplat}.flags = zeroSize ? 0u : GSPLAT_FLAG_ACTIVE;
|
||||
} else {
|
||||
${gsplat}.flags = 0u;
|
||||
}
|
||||
`);
|
||||
} else {
|
||||
statements = [`${gsplat}.flags = 0u;`];
|
||||
}
|
||||
statements.push(`${gsplat}.index = ${index ?? "0"};`);
|
||||
return statements;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
dynoOut(): DynoValue<typeof Gsplat> {
|
||||
return new DynoOutput(this, "gsplat");
|
||||
}
|
||||
}
|
||||
|
||||
export class ReadPackedSplatRange
|
||||
extends Dyno<
|
||||
{
|
||||
packedSplats: typeof TPackedSplats;
|
||||
index: "int";
|
||||
base: "int";
|
||||
count: "int";
|
||||
},
|
||||
{ gsplat: typeof Gsplat }
|
||||
>
|
||||
implements HasDynoOut<typeof Gsplat>
|
||||
{
|
||||
constructor({
|
||||
packedSplats,
|
||||
index,
|
||||
base,
|
||||
count,
|
||||
}: {
|
||||
packedSplats?: DynoVal<typeof TPackedSplats>;
|
||||
index?: DynoVal<"int">;
|
||||
base?: DynoVal<"int">;
|
||||
count?: DynoVal<"int">;
|
||||
}) {
|
||||
super({
|
||||
inTypes: {
|
||||
packedSplats: TPackedSplats,
|
||||
index: "int",
|
||||
base: "int",
|
||||
count: "int",
|
||||
},
|
||||
outTypes: { gsplat: Gsplat },
|
||||
inputs: { packedSplats, index, base, count },
|
||||
globals: () => [defineGsplat, definePackedSplats, defineReadPackedSplat],
|
||||
statements: ({ inputs, outputs }) => {
|
||||
const { gsplat } = outputs;
|
||||
if (!gsplat) {
|
||||
return [];
|
||||
}
|
||||
const { packedSplats, index, base, count } = inputs;
|
||||
let statements: string[];
|
||||
if (packedSplats && index && base && count) {
|
||||
statements = unindentLines(`
|
||||
${gsplat}.flags = 0u;
|
||||
if ((${index} >= ${base}) && (${index} < (${base} + ${count}))) {
|
||||
if (readPackedSplat(${packedSplats}.texture, ${packedSplats}.numSplats, ${index}, ${gsplat})) {
|
||||
bool zeroSize = all(equal(${gsplat}.scales, vec3(0.0, 0.0, 0.0)));
|
||||
${gsplat}.flags = zeroSize ? 0u : GSPLAT_FLAG_ACTIVE;
|
||||
}
|
||||
}
|
||||
`);
|
||||
} else {
|
||||
statements = [`${gsplat}.flags = 0u;`];
|
||||
}
|
||||
statements.push(`${gsplat}.index = ${index ?? "0"};`);
|
||||
return statements;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
dynoOut(): DynoValue<typeof Gsplat> {
|
||||
return new DynoOutput(this, "gsplat");
|
||||
}
|
||||
}
|
||||
|
||||
export class SplitGsplat extends Dyno<
|
||||
{ gsplat: typeof Gsplat },
|
||||
{
|
||||
flags: "uint";
|
||||
active: "bool";
|
||||
index: "int";
|
||||
center: "vec3";
|
||||
scales: "vec3";
|
||||
quaternion: "vec4";
|
||||
rgba: "vec4";
|
||||
rgb: "vec3";
|
||||
opacity: "float";
|
||||
x: "float";
|
||||
y: "float";
|
||||
z: "float";
|
||||
r: "float";
|
||||
g: "float";
|
||||
b: "float";
|
||||
}
|
||||
> {
|
||||
constructor({ gsplat }: { gsplat?: DynoVal<typeof Gsplat> }) {
|
||||
super({
|
||||
inTypes: { gsplat: Gsplat },
|
||||
outTypes: {
|
||||
flags: "uint",
|
||||
active: "bool",
|
||||
index: "int",
|
||||
center: "vec3",
|
||||
scales: "vec3",
|
||||
quaternion: "vec4",
|
||||
rgba: "vec4",
|
||||
rgb: "vec3",
|
||||
opacity: "float",
|
||||
x: "float",
|
||||
y: "float",
|
||||
z: "float",
|
||||
r: "float",
|
||||
g: "float",
|
||||
b: "float",
|
||||
},
|
||||
inputs: { gsplat },
|
||||
globals: () => [defineGsplat],
|
||||
statements: ({ inputs, outputs }) => {
|
||||
const { gsplat } = inputs;
|
||||
const {
|
||||
flags,
|
||||
active,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
} = outputs;
|
||||
return [
|
||||
!flags ? null : `${flags} = ${gsplat ? `${gsplat}.flags` : "0u"};`,
|
||||
!active
|
||||
? null
|
||||
: `${active} = isGsplatActive(${gsplat ? `${gsplat}.flags` : "0u"});`,
|
||||
!index ? null : `${index} = ${gsplat ? `${gsplat}.index` : "0"};`,
|
||||
!center
|
||||
? null
|
||||
: `${center} = ${gsplat ? `${gsplat}.center` : "vec3(0.0, 0.0, 0.0)"};`,
|
||||
!scales
|
||||
? null
|
||||
: `${scales} = ${gsplat ? `${gsplat}.scales` : "vec3(0.0, 0.0, 0.0)"};`,
|
||||
!quaternion
|
||||
? null
|
||||
: `${quaternion} = ${gsplat ? `${gsplat}.quaternion` : "vec4(0.0, 0.0, 0.0, 1.0)"};`,
|
||||
!rgba
|
||||
? null
|
||||
: `${rgba} = ${gsplat ? `${gsplat}.rgba` : "vec4(0.0, 0.0, 0.0, 0.0)"};`,
|
||||
!rgb
|
||||
? null
|
||||
: `${rgb} = ${gsplat ? `${gsplat}.rgba.rgb` : "vec3(0.0, 0.0, 0.0)"};`,
|
||||
!opacity
|
||||
? null
|
||||
: `${opacity} = ${gsplat ? `${gsplat}.rgba.a` : "0.0"};`,
|
||||
!x ? null : `${x} = ${gsplat ? `${gsplat}.center.x` : "0.0"};`,
|
||||
!y ? null : `${y} = ${gsplat ? `${gsplat}.center.y` : "0.0"};`,
|
||||
!z ? null : `${z} = ${gsplat ? `${gsplat}.center.z` : "0.0"};`,
|
||||
!r ? null : `${r} = ${gsplat ? `${gsplat}.rgba.r` : "0.0"};`,
|
||||
!g ? null : `${g} = ${gsplat ? `${gsplat}.rgba.g` : "0.0"};`,
|
||||
!b ? null : `${b} = ${gsplat ? `${gsplat}.rgba.b` : "0.0"};`,
|
||||
].filter(Boolean) as string[];
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export class CombineGsplat
|
||||
extends Dyno<
|
||||
{
|
||||
gsplat: typeof Gsplat;
|
||||
flags: "uint";
|
||||
index: "int";
|
||||
center: "vec3";
|
||||
scales: "vec3";
|
||||
quaternion: "vec4";
|
||||
rgba: "vec4";
|
||||
rgb: "vec3";
|
||||
opacity: "float";
|
||||
x: "float";
|
||||
y: "float";
|
||||
z: "float";
|
||||
r: "float";
|
||||
g: "float";
|
||||
b: "float";
|
||||
},
|
||||
{ gsplat: typeof Gsplat }
|
||||
>
|
||||
implements HasDynoOut<typeof Gsplat>
|
||||
{
|
||||
constructor({
|
||||
gsplat,
|
||||
flags,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
}: {
|
||||
gsplat?: DynoVal<typeof Gsplat>;
|
||||
flags?: DynoVal<"uint">;
|
||||
index?: DynoVal<"int">;
|
||||
center?: DynoVal<"vec3">;
|
||||
scales?: DynoVal<"vec3">;
|
||||
quaternion?: DynoVal<"vec4">;
|
||||
rgba?: DynoVal<"vec4">;
|
||||
rgb?: DynoVal<"vec3">;
|
||||
opacity?: DynoVal<"float">;
|
||||
x?: DynoVal<"float">;
|
||||
y?: DynoVal<"float">;
|
||||
z?: DynoVal<"float">;
|
||||
r?: DynoVal<"float">;
|
||||
g?: DynoVal<"float">;
|
||||
b?: DynoVal<"float">;
|
||||
}) {
|
||||
super({
|
||||
inTypes: {
|
||||
gsplat: Gsplat,
|
||||
flags: "uint",
|
||||
index: "int",
|
||||
center: "vec3",
|
||||
scales: "vec3",
|
||||
quaternion: "vec4",
|
||||
rgba: "vec4",
|
||||
rgb: "vec3",
|
||||
opacity: "float",
|
||||
x: "float",
|
||||
y: "float",
|
||||
z: "float",
|
||||
r: "float",
|
||||
g: "float",
|
||||
b: "float",
|
||||
},
|
||||
outTypes: { gsplat: Gsplat },
|
||||
inputs: {
|
||||
gsplat,
|
||||
flags,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
},
|
||||
globals: () => [defineGsplat],
|
||||
statements: ({ inputs, outputs }) => {
|
||||
const { gsplat: outGsplat } = outputs;
|
||||
if (!outGsplat) {
|
||||
return [];
|
||||
}
|
||||
const {
|
||||
gsplat,
|
||||
flags,
|
||||
index,
|
||||
center,
|
||||
scales,
|
||||
quaternion,
|
||||
rgba,
|
||||
rgb,
|
||||
opacity,
|
||||
x,
|
||||
y,
|
||||
z,
|
||||
r,
|
||||
g,
|
||||
b,
|
||||
} = inputs;
|
||||
return [
|
||||
`${outGsplat}.flags = ${flags ?? (gsplat ? `${gsplat}.flags` : "0u")};`,
|
||||
`${outGsplat}.index = ${index ?? (gsplat ? `${gsplat}.index` : "0")};`,
|
||||
`${outGsplat}.center = ${center ?? (gsplat ? `${gsplat}.center` : "vec3(0.0, 0.0, 0.0)")};`,
|
||||
`${outGsplat}.scales = ${scales ?? (gsplat ? `${gsplat}.scales` : "vec3(0.0, 0.0, 0.0)")};`,
|
||||
`${outGsplat}.quaternion = ${quaternion ?? (gsplat ? `${gsplat}.quaternion` : "vec4(0.0, 0.0, 0.0, 1.0)")};`,
|
||||
`${outGsplat}.rgba = ${rgba ?? (gsplat ? `${gsplat}.rgba` : "vec4(0.0, 0.0, 0.0, 0.0)")};`,
|
||||
!rgb ? null : `${outGsplat}.rgba.rgb = ${rgb};`,
|
||||
!opacity ? null : `${outGsplat}.rgba.a = ${opacity};`,
|
||||
!x ? null : `${outGsplat}.center.x = ${x};`,
|
||||
!y ? null : `${outGsplat}.center.y = ${y};`,
|
||||
!z ? null : `${outGsplat}.center.z = ${z};`,
|
||||
!r ? null : `${outGsplat}.rgba.r = ${r};`,
|
||||
!g ? null : `${outGsplat}.rgba.g = ${g};`,
|
||||
!b ? null : `${outGsplat}.rgba.b = ${b};`,
|
||||
].filter(Boolean) as string[];
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
dynoOut(): DynoValue<typeof Gsplat> {
|
||||
return new DynoOutput(this, "gsplat");
|
||||
}
|
||||
}
|
||||
|
||||
export const defineGsplatNormal = unindent(`
|
||||
vec3 gsplatNormal(vec3 scales, vec4 quaternion) {
|
||||
float minScale = min(scales.x, min(scales.y, scales.z));
|
||||
vec3 normal;
|
||||
if (scales.z == minScale) {
|
||||
normal = vec3(0.0, 0.0, 1.0);
|
||||
} else if (scales.y == minScale) {
|
||||
normal = vec3(0.0, 1.0, 0.0);
|
||||
} else {
|
||||
normal = vec3(1.0, 0.0, 0.0);
|
||||
}
|
||||
return quatVec(quaternion, normal);
|
||||
}
|
||||
`);
|
||||
|
||||
export class GsplatNormal extends UnaryOp<typeof Gsplat, "vec3", "normal"> {
|
||||
constructor({ gsplat }: { gsplat: DynoVal<typeof Gsplat> }) {
|
||||
super({ a: gsplat, outKey: "normal", outTypeFunc: () => "vec3" });
|
||||
this.globals = () => [defineGsplat, defineGsplatNormal];
|
||||
this.statements = ({ inputs, outputs }) => [
|
||||
`${outputs.normal} = gsplatNormal(${inputs.a}.scales, ${inputs.a}.quaternion);`,
|
||||
];
|
||||
}
|
||||
}
|
||||
|
||||
export class TransformGsplat
|
||||
extends Dyno<
|
||||
{
|
||||
gsplat: typeof Gsplat;
|
||||
scale: "float";
|
||||
rotate: "vec4";
|
||||
translate: "vec3";
|
||||
recolor: "vec4";
|
||||
},
|
||||
{ gsplat: typeof Gsplat }
|
||||
>
|
||||
implements HasDynoOut<typeof Gsplat>
|
||||
{
|
||||
constructor({
|
||||
gsplat,
|
||||
scale,
|
||||
rotate,
|
||||
translate,
|
||||
recolor,
|
||||
}: {
|
||||
gsplat?: DynoVal<typeof Gsplat>;
|
||||
scale?: DynoVal<"float">;
|
||||
rotate?: DynoVal<"vec4">;
|
||||
translate?: DynoVal<"vec3">;
|
||||
recolor?: DynoVal<"vec4">;
|
||||
}) {
|
||||
super({
|
||||
inTypes: {
|
||||
gsplat: Gsplat,
|
||||
scale: "float",
|
||||
rotate: "vec4",
|
||||
translate: "vec3",
|
||||
recolor: "vec4",
|
||||
},
|
||||
outTypes: { gsplat: Gsplat },
|
||||
inputs: { gsplat, scale, rotate, translate, recolor },
|
||||
globals: () => [defineGsplat],
|
||||
statements: ({ inputs, outputs, compile }) => {
|
||||
const { gsplat } = outputs;
|
||||
if (!gsplat || !inputs.gsplat) {
|
||||
return [];
|
||||
}
|
||||
const { scale, rotate, translate, recolor } = inputs;
|
||||
const indent = compile.indent;
|
||||
const statements = [
|
||||
`${gsplat} = ${inputs.gsplat};`,
|
||||
`if (isGsplatActive(${gsplat}.flags)) {`,
|
||||
|
||||
scale ? `${indent}${gsplat}.center *= ${scale};` : null,
|
||||
rotate
|
||||
? `${indent}${gsplat}.center = quatVec(${rotate}, ${gsplat}.center);`
|
||||
: null,
|
||||
translate ? `${indent}${gsplat}.center += ${translate};` : null,
|
||||
|
||||
scale ? `${indent}${gsplat}.scales *= ${scale};` : null,
|
||||
|
||||
rotate
|
||||
? `${indent}${gsplat}.quaternion = quatQuat(${rotate}, ${gsplat}.quaternion);`
|
||||
: null,
|
||||
recolor ? `${indent}${gsplat}.rgba *= ${recolor};` : null,
|
||||
"}",
|
||||
].filter(Boolean) as string[];
|
||||
return statements;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
dynoOut(): DynoValue<typeof Gsplat> {
|
||||
return new DynoOutput(this, "gsplat");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user