glre 0.30.0 → 0.32.0

package/src/node/code.ts

@@ -1,11 +1,32 @@
 import { is } from '../utils/helpers'
 import { infer } from './infer'
-import { getBluiltin, getOperator, generateDefine, formatConversions, joins } from './utils'
-import type { NodeConfig, X } from './types'
+import { getBluiltin, getOperator, formatConversions, joins } from './utils'
+import type { NodeContext, X } from './types'
+import {
+        parseAttribHead,
+        parseConstantHead,
+        parseUniformHead,
+        parseDeclare,
+        parseDefine,
+        parseIf,
+        parseSwitch,
+        parseTexture,
+        parseVaryingHead,
+} from './parse'
 
-export const code = (target: X, c?: NodeConfig | null): string => {
+export const code = (target: X, c?: NodeContext | null): string => {
         if (!c) c = {}
         if (!c.headers) c.headers = new Map()
+        if (!c.vertVaryings) c.vertVaryings = new Map()
+        if (!c.fragInputs) c.fragInputs = new Map()
+        if (!c.vertInputs) c.vertInputs = new Map()
+        if (!c.vertOutputs) {
+                c.vertOutputs = new Map()
+                if (!c.isWebGL) {
+                        c.fragInputs.set('position', '@builtin(position) position: vec4f')
+                        c.vertOutputs.set('position', '@builtin(position) position: vec4f')
+                }
+        }
         if (is.str(target)) return target
         if (is.num(target)) {
                 const ret = `${target}`
@@ -13,50 +34,19 @@ export const code = (target: X, c?: NodeConfig | null): string => {
                 return ret + '.0'
         }
         if (is.bol(target)) return target ? 'true' : 'false'
-        if (!target) return '' // ignore if no target
+        if (!target) return ''
         const { type, props } = target
         const { id = '', children = [] } = props
-        const [x, y, z] = children
-        /**
-         * headers
-         */
-        let head = ''
-        if (type === 'attribute') {
-                if (c.headers.has(id)) return id
-                head = `${infer(target, c)} ${id}`
-        }
-        if (type === 'uniform') {
-                if (c.headers.has(id)) return id
-                if (!c.binding) c.binding = 0
-                const varType = infer(target, c)
-                head = c.isWebGL
-                        ? `uniform ${varType} ${id};`
-                        : `@group(0) @binding(${c.binding++}) var<uniform> ${id}: ${formatConversions(varType, c)};`
-        }
-        if (type === 'constant') {
-                if (c.headers.has(id)) return id
-                const varType = infer(target, c)
-                const value = code(x, c)
-                head = c.isWebGL
-                        ? `const ${varType} ${id} = ${value};`
-                        : `const ${id}: ${formatConversions(varType, c)} = ${value};`
-        }
-        if (type === 'varying') {
-                if (c.headers.has(id)) return id
-                head = `${infer(target, c)} ${id}`
-        }
-        if (head) {
-                c.headers.set(id, head)
-                c.onMount?.(id)
-                return id
-        }
+        const [x, y, z, w] = children
         /**
          * variables
          */
         if (type === 'variable') return id
         if (type === 'swizzle') return `${code(y, c)}.${code(x, c)}`
-        if (type === 'ternary') return `(${code(x, c)} ? ${code(y, c)} : ${code(z, c)})`
-        if (type === 'builtin') return c?.isWebGL ? getBluiltin(id) : id
+        if (type === 'ternary')
+                return c.isWebGL
+                        ? `(${code(x, c)} ? ${code(y, c)} : ${code(z, c)})`
+                        : `select(${code(x, c)}, ${code(y, c)}, ${code(z, c)})`
         if (type === 'conversion') return `${formatConversions(x, c)}(${joins(children.slice(1), c)})`
         if (type === 'operator') {
                 if (x === 'not' || x === 'bitNot') return `!${code(y, c)}`
@@ -64,6 +54,7 @@ export const code = (target: X, c?: NodeConfig | null): string => {
         }
         if (type === 'function') {
                 if (x === 'negate') return `(-${joins(children.slice(1), c)})`
+                if (x === 'texture') return parseTexture(c, y, z, w)
                 return `${x}(${joins(children.slice(1), c)})`
         }
         /**
@@ -71,45 +62,51 @@ export const code = (target: X, c?: NodeConfig | null): string => {
          */
         if (type === 'scope') return children.map((child: any) => code(child, c)).join('\n')
         if (type === 'assign') return `${code(x, c)} = ${code(y, c)};`
+        if (type === 'return') return `return ${code(x, c)};`
         if (type === 'loop')
                 return c.isWebGL
                         ? `for (int i = 0; i < ${x}; i += 1) {\n${code(y, c)}\n}`
                         : `for (var i: i32 = 0; i < ${x}; i++) {\n${code(y, c)}\n}`
+        if (type === 'if') return parseIf(c, x, y, children)
+        if (type === 'switch') return parseSwitch(c, x, children)
+        if (type === 'declare') return parseDeclare(c, x, y)
         if (type === 'define') {
-                const args = children.slice(2)
-                const ret = `${id}(${args.map((arg) => code(arg, c))})`
+                const ret = `${id}(${joins(children.slice(1), c)})`
                 if (c.headers.has(id)) return ret
-                c.headers.set(id, generateDefine(props, c))
+                c.headers.set(id, parseDefine(c, props, infer(target, c)))
                 return ret
         }
-        if (type === 'if') {
-                let ret = `if (${code(x, c)}) {\n${code(y, c)}\n}`
-                for (let i = 2; i < children.length; i += 2) {
-                        const isElseIf = i >= children.length - 1
-                        ret += !isElseIf
-                                ? ` else if (${code(children[i], c)}) {\n${code(children[i + 1], c)}\n}`
-                                : ` else {\n${code(children[i], c)}\n}`
-                }
-                return ret
+        /**
+         * headers
+         */
+        if (type === 'varying') {
+                if (c.vertOutputs.has(id)) return c.isWebGL ? `${id}` : `out.${id}`
+                const field = parseVaryingHead(c, id, infer(target, c))
+                c.fragInputs.set(id, field)
+                c.vertOutputs.set(id, field)
+                c.vertVaryings.set(id, code(x, c))
+                return c.isWebGL ? `${id}` : `out.${id}`
         }
-        if (type === 'switch') {
-                let ret = `switch (${code(x, c)}) {\n`
-                for (let i = 1; i < children.length; i += 2) {
-                        const isDefault = i >= children.length - 1
-                        if (isDefault && children.length % 2 === 0) {
-                                ret += `default:\n${code(children[i], c)}\nbreak;\n`
-                        } else if (i + 1 < children.length)
-                                ret += `case ${code(children[i], c)}:\n${code(children[i + 1], c)}\nbreak;\n`
-                }
-                ret += '}'
-                return ret
+        if (type === 'builtin') {
+                if (c.isWebGL) return getBluiltin(id)
+                if (id === 'position') return 'out.position'
+                const field = `@builtin(${id}) ${id}: ${formatConversions(infer(target, c), c)}`
+                if (c.isFrag) {
+                        c.fragInputs.set(id, field)
+                } else c.vertInputs.set(id, field)
+                return `in.${id}`
+        }
+        if (type === 'attribute') {
+                c.vertInputs.set(id, parseAttribHead(c, id, infer(target, c)))
+                return c.isWebGL ? `${id}` : `in.${id}`
         }
-        if (type === 'declare') {
-                const varType = infer(x, c)
-                const varName = (y as any)?.props?.id
-                if (c.isWebGL) return `${varType} ${varName} = ${code(x, c)};`
-                const wgslType = formatConversions(varType)
-                return `var ${varName}: ${wgslType} = ${code(x, c)};`
+        if (c.headers.has(id)) return id // must last
+        let head = ''
+        if (type === 'uniform') head = parseUniformHead(c, id, infer(target, c))
+        if (type === 'constant') head = parseConstantHead(c, id, infer(target, c), code(x, c))
+        if (head) {
+                c.headers.set(id, head)
+                return id
         }
         return code(x, c)
 }

package/src/node/const.ts

@@ -21,6 +21,8 @@ export const CONSTANTS = [
         'mat2',
         'mat3',
         'mat4',
+        'texture',
+        'sampler2D',
 ] as const
 
 export const CONVERSIONS = [

package/src/node/index.ts

@@ -1,15 +1,82 @@
-import { builtin, conversion as c, function_ as f, uniform } from './node'
-import { hex2rgb } from './utils'
 import { is } from '../utils/helpers'
-import type { X } from './types'
+import { code } from './code'
+import { builtin, conversion as c, function_ as f, uniform as u } from './node'
+import { hex2rgb } from './utils'
+import type { NodeContext, X } from './types'
 export * from './code'
-export * from './const'
-export * from './infer'
 export * from './node'
 export * from './scope'
 export * from './types'
 export * from './utils'
 
+const GLSL_FRAGMENT_HEAD = `
+#version 300 es
+precision mediump float;
+out vec4 fragColor;
+`.trim()
+
+const generateHead = (x: X, c: NodeContext) => {
+        const body = code(x, c)
+        const head = Array.from(c.headers!.values()).join('\n')
+        return [head, body]
+}
+
+const generateStruct = (id: string, map: Map<string, string>) => {
+        return `struct ${id} {\n  ${Array.from(map.values()).join(',\n  ')}\n}`
+}
+
+export const vertex = (x: X, c: NodeContext) => {
+        c.headers?.clear()
+        c.isFrag = false // for varying inputs or outputs
+        const [head, body] = generateHead(x, c)
+        const ret = []
+        if (c.isWebGL) {
+                ret.push('#version 300 es')
+                for (const code of c.vertInputs!.values()) ret.push(`in ${code}`)
+                for (const code of c.vertOutputs!.values()) ret.push(`out ${code}`)
+                ret.push(head)
+                ret.push('void main() {')
+                ret.push(`  gl_Position = ${body};`)
+                for (const [id, code] of c.vertVaryings!.entries()) ret.push(`  ${id} = ${code};`)
+        } else {
+                if (c.vertInputs?.size) ret.push(generateStruct('In', c.vertInputs))
+                if (c.vertOutputs?.size) ret.push(generateStruct('Out', c.vertOutputs))
+                ret.push(head)
+                ret.push('@vertex')
+                ret.push(`fn main(${c.vertInputs?.size ? 'in: In' : ''}) -> Out {`)
+                ret.push('  var out: Out;')
+                ret.push(`  out.position = ${body};`)
+                for (const [id, code] of c.vertVaryings!.entries()) ret.push(`  out.${id} = ${code};`)
+                ret.push('  return out;')
+        }
+        ret.push('}')
+        const main = ret.filter(Boolean).join('\n')
+        console.log(`↓↓↓generated↓↓↓\n${main}`)
+        return main
+}
+
+export const fragment = (x: X, c: NodeContext) => {
+        c.headers?.clear()
+        c.isFrag = true // for varying inputs or outputs
+        const [head, body] = generateHead(x, c)
+        const ret = []
+        if (c.isWebGL) {
+                ret.push(GLSL_FRAGMENT_HEAD)
+                for (const code of c.fragInputs!.values()) ret.push(`in ${code}`)
+                ret.push(head)
+                ret.push(`void main() {\n  fragColor = ${body};`)
+        } else {
+                if (c.fragInputs?.size) ret.push(generateStruct('Out', c.fragInputs))
+                ret.push(head)
+                ret.push(`@fragment\nfn main(out: Out) -> @location(0) vec4f {`)
+                ret.push(`  return ${body};`)
+        }
+        ret.push('}')
+        const main = ret.filter(Boolean).join('\n')
+        console.log(`↓↓↓generated↓↓↓\n${main}`)
+        return main
+}
+
 // Builtin Variables
 export const position = builtin('position')
 export const vertexIndex = builtin('vertex_index')
@@ -50,15 +117,18 @@ export const uvec4 = (x?: X, y?: X, z?: X, w?: X) => c('uvec4', x, y, z, w)
 export const bvec2 = (x?: X, y?: X) => c('bvec2', x, y)
 export const bvec3 = (x?: X, y?: X, z?: X) => c('bvec3', x, y, z)
 export const bvec4 = (x?: X, y?: X, z?: X, w?: X) => c('bvec4', x, y, z, w)
+export const texture2D = () => c('texture')
+export const sampler2D = () => c('sampler2D')
 export const color = (r?: X, g?: X, b?: X) => {
         if (is.num(r) && is.und(g) && is.und(b)) return vec3(...hex2rgb(r))
         return vec3(r, g, b)
 }
 
 // Default uniforms
-export const iResolution = uniform(vec2(1280, 800), 'iResolution')
-export const iMouse = uniform(vec2(0, 0), 'iMouse')
-export const iTime = uniform(float(0), 'iTime')
+export const iResolution = u(vec2(1280, 800), 'iResolution')
+export const iMouse = u(vec2(0, 0), 'iMouse')
+export const iTime = u(float(0), 'iTime')
+export const uv = () => position.xy.div(iResolution)
 
 // Texture Functions
 export const texture = (x: X, y: X, z?: X) => f('texture', x, y, z)

package/src/node/infer.ts

@@ -13,10 +13,10 @@ import {
         VEC3_RETURN_FUNCTIONS,
         VEC4_RETURN_FUNCTIONS,
 } from './const'
-import { isNodeProxy } from './utils'
-import type { Constants, NodeConfig, NodeProxy, X } from './types'
+import { isConstantsType, isNodeProxy } from './utils'
+import type { Constants, NodeContext, NodeProxy, X } from './types'
 
-const getHighestPriorityType = (args: X[], c: NodeConfig) => {
+const getHighestPriorityType = (args: X[], c: NodeContext) => {
         return args.reduce((highest, current) => {
                 const currentType = infer(current, c)
                 const highestPriority = CONSTANTS.indexOf(highest as any)
@@ -29,7 +29,7 @@ const inferBuiltin = (id: string | undefined): Constants => {
         return BUILTIN_TYPES[id as keyof typeof BUILTIN_TYPES]!
 }
 
-const inferFunction = (funcName: string, args: X[], c: NodeConfig): Constants => {
+const inferFunction = (funcName: string, args: X[], c: NodeContext): Constants => {
         const firstArgType = args.length > 0 ? infer(args[0], c) : 'float'
         if (FIRST_ARG_TYPE_FUNCTIONS.includes(funcName as any)) return firstArgType
         if (SCALAR_RETURN_FUNCTIONS.includes(funcName as any)) return 'float'
@@ -52,9 +52,9 @@ const inferOperator = (leftType: string, rightType: string, op: string): Constan
         return (leftPriority >= rightPriority ? leftType : rightType) as Constants
 }
 
-const inferPrimitiveType = (x: any): Constants => {
+export const inferPrimitiveType = (x: any): Constants => {
         if (is.bol(x)) return 'bool'
-        if (is.num(x)) return 'float' // Number.isInteger(x) ? 'int' : 'float' // @TODO FIX
+        if (is.num(x)) return 'float'
         if (is.arr(x)) return COMPONENT_COUNT_TO_TYPE[x.length as keyof typeof COMPONENT_COUNT_TO_TYPE] || 'float'
         return 'float'
 }
@@ -63,25 +63,30 @@ const inferSwizzle = (count: number): Constants => {
         return COMPONENT_COUNT_TO_TYPE[count as keyof typeof COMPONENT_COUNT_TO_TYPE]!
 }
 
-export const inferImpl = (target: NodeProxy, c: NodeConfig): Constants => {
+const inferFromArray = (arr: X[], c: NodeContext): Constants => {
+        if (arr.length === 0) return 'void'
+        const ret = infer(arr[0], c)
+        for (const x of arr.slice(1))
+                if (ret !== infer(x)) throw new Error(`glre node system error: defined scope return mismatch`)
+        return ret
+}
+
+export const inferImpl = (target: NodeProxy, c: NodeContext): Constants => {
         const { type, props } = target
-        const { id, children = [], inferFrom } = props
+        const { id, children = [], layout, inferFrom } = props
         const [x, y, z] = children
-        if (inferFrom) return infer(inferFrom, c)
         if (type === 'conversion') return x as Constants
         if (type === 'operator') return inferOperator(infer(y, c), infer(z, c), x as string)
         if (type === 'function') return inferFunction(x as string, children.slice(1), c)
         if (type === 'swizzle') return inferSwizzle((x as string).length)
         if (type === 'ternary') return inferOperator(infer(y, c), infer(z, c), 'add')
         if (type === 'builtin') return inferBuiltin(id)
-        if (type === 'define') {
-                if (props.layout?.type && props.layout?.type !== 'auto') return props.layout.type as Constants
-                return y ? infer(y, c) : 'void'
-        }
+        if (type === 'define' && isConstantsType(layout?.type)) return layout?.type
+        if (inferFrom) return inferFromArray(inferFrom, c)
         return infer(x, c)
 }
 
-export const infer = (target: X, c?: NodeConfig | null): Constants => {
+export const infer = (target: X, c?: NodeContext | null): Constants => {
         if (!c) c = {}
         if (!isNodeProxy(target)) return inferPrimitiveType(target)
         if (!c.infers) c.infers = new WeakMap<NodeProxy, Constants>()

package/src/node/node.ts

@@ -1,6 +1,6 @@
 import { code } from './code'
 import { assign, toVar } from './scope'
-import { conversionToConstant, isConversion, isFunction, isOperator, isSwizzle } from './utils'
+import { conversionToConstant, isConversion, isFunction, isOperator, isSwizzle, getId, isNodeProxy } from './utils'
 import type { Functions, NodeProps, NodeProxy, NodeTypes, Operators, Swizzles, X } from './types'
 
 const toPrimitive = (x: X, hint: string) => {
@@ -35,18 +35,16 @@ export const node = (type: NodeTypes, props?: NodeProps | null, ...args: X[]) =>
 }
 
 // headers
-export const attribute = (x: X, id?: string) => node('varying', { id }, x)
-export const uniform = (x: X, id?: string) => node('uniform', { id }, x)
-export const varying = (x: X, id?: string) => node('varying', { id }, x)
-export const constant = (x: X, id?: string) => node('constant', { id }, x)
-export const variable = (id: string) => node('variable', { id })
-export const builtin = (id: string) => node('builtin', { id })
+export const attribute = (x: X, id = getId()) => node('attribute', { id }, x)
+export const constant = (x: X, id = getId()) => node('constant', { id }, x)
+export const uniform = (x: X, id = getId()) => node('uniform', { id }, x)
+export const variable = (id = getId()) => node('variable', { id })
+export const builtin = (id = getId()) => node('builtin', { id })
+export const vertexStage = (x: X, id = getId()) => node('varying', { id, inferFrom: [x] }, x)
 
 // Node shorthands
-export const swizzle = (key: Swizzles, arg: X) => node('swizzle', null, key, arg)
-export const operator = (key: Operators, ...args: X[]) => node('operator', null, key, ...args)
-export const function_ = (key: Functions, ...args: X[]) => node('function', null, key, ...args)
-export const conversion = (key: string, ...args: X[]) => node('conversion', null, key, ...args)
-
-// x ? y : z
-export const select = (x: X, y: X, z: X) => node('ternary', null, x, y, z)
+export const swizzle = (key: Swizzles, x: X) => node('swizzle', null, key, x)
+export const operator = (key: Operators, ...x: X[]) => node('operator', null, key, ...x)
+export const function_ = (key: Functions, ...x: X[]) => node('function', null, key, ...x)
+export const conversion = (key: string, ...x: X[]) => node('conversion', null, key, ...x)
+export const select = (x: X, y: X, z: X) => node('ternary', null, x, y, z) // x ? y : z

package/src/node/parse.ts

@@ -0,0 +1,118 @@
+import { infer } from './infer'
+import { formatConversions, joins } from './utils'
+import { code } from './code'
+import type { Constants, NodeContext, NodeProps, X } from './types'
+
+export const parseTexture = (c: NodeContext, y: X, z: X, w: X) => {
+        if (c.isWebGL) {
+                const args = w ? [y, z, w] : [y, z]
+                return `texture(${joins(args, c)})`
+        }
+        const _y = code(y, c)
+        const args = [_y, _y + 'Sampler', code(z, c)]
+        if (!w) return `textureSample(${args})`
+        args.push(code(w, c))
+        return `textureSampleLevel(${args})`
+}
+
+/**
+ * scopes
+ */
+export const parseIf = (c: NodeContext, x: X, y: X, children: X[]) => {
+        let ret = `if (${code(x, c)}) {\n${code(y, c)}\n}`
+        for (let i = 2; i < children.length; i += 2) {
+                const isElse = i >= children.length - 1
+                ret += !isElse
+                        ? ` else if (${code(children[i], c)}) {\n${code(children[i + 1], c)}\n}`
+                        : ` else {\n${code(children[i], c)}\n}`
+        }
+        return ret
+}
+
+export const parseSwitch = (c: NodeContext, x: X, children: X[]) => {
+        let ret = `switch (${code(x, c)}) {\n`
+        for (let i = 1; i < children.length; i += 2) {
+                const isDefault = i >= children.length - 1
+                if (isDefault && children.length % 2 === 0) {
+                        ret += `default:\n${code(children[i], c)}\nbreak;\n`
+                } else if (i + 1 < children.length)
+                        ret += `case ${code(children[i], c)}:\n${code(children[i + 1], c)}\nbreak;\n`
+        }
+        ret += '}'
+        return ret
+}
+
+export const parseDeclare = (c: NodeContext, x: X, y: X) => {
+        const varType = infer(x, c)
+        const varName = (y as any)?.props?.id
+        if (c.isWebGL) return `${varType} ${varName} = ${code(x, c)};`
+        const wgslType = formatConversions(varType)
+        return `var ${varName}: ${wgslType} = ${code(x, c)};`
+}
+
+export const parseDefine = (c: NodeContext, props: NodeProps, returnType: Constants) => {
+        const { id, children = [], layout } = props
+        const [x, ...args] = children
+        const argParams: [name: string, type: string][] = []
+        const params: string[] = []
+        if (layout?.inputs)
+                for (const input of layout.inputs) {
+                        argParams.push([input.name, input.type])
+                }
+        else
+                for (let i = 0; i < args.length; i++) {
+                        argParams.push([`p${i}`, infer(args[i], c)])
+                }
+        const ret = []
+        if (c?.isWebGL) {
+                for (const [id, type] of argParams) params.push(`${type} ${id}`)
+                ret.push(`${returnType} ${id}(${params}) {`)
+        } else {
+                for (const [id, type] of argParams) params.push(`${id}: ${formatConversions(type, c)}`)
+                ret.push(`fn ${id}(${params}) -> ${formatConversions(returnType, c)} {`)
+        }
+        const scopeCode = code(x, c)
+        if (scopeCode) ret.push(scopeCode)
+        ret.push('}')
+        return ret.join('\n')
+}
+
+/**
+ * headers
+ */
+export const parseVaryingHead = (c: NodeContext, id: string, varType: string) => {
+        return c.isWebGL
+                ? `${varType} ${id};`
+                : `@location(${c.vertVaryings!.size}) ${id}: ${formatConversions(varType, c)}`
+}
+
+export const parseUniformHead = (c: NodeContext, id: string, varType: Constants) => {
+        const isTexture = varType === 'sampler2D' || varType === 'texture'
+        if (c.isWebGL)
+                return isTexture //
+                        ? `uniform sampler2D ${id};`
+                        : `uniform ${varType} ${id};`
+        if (isTexture) {
+                const { group = 1, binding = 0 } = c.webgpu?.textures.map.get(id) || {}
+                return (
+                        `@group(${group}) @binding(${binding}) var ${id}Sampler: sampler;\n` +
+                        `@group(${group}) @binding(${binding + 1}) var ${id}: texture_2d<f32>;`
+                )
+        }
+        const { group = 0, binding = 0 } = c.webgpu?.uniforms.map.get(id) || {}
+        const wgslType = formatConversions(varType, c)
+        return `@group(${group}) @binding(${binding}) var<uniform> ${id}: ${wgslType};`
+}
+
+export const parseAttribHead = (c: NodeContext, id: string, varType: Constants) => {
+        if (c.isWebGL) return `${varType} ${id};`
+        const { location = 0 } = c.webgpu?.attribs.map.get(id) || {}
+        const wgslType = formatConversions(varType, c)
+        return `@location(${location}) ${id}: ${wgslType}`
+}
+
+export const parseConstantHead = (c: NodeContext, id: string, varType: Constants, value: string) => {
+        return c.isWebGL
+                ? `const ${varType} ${id} = ${value};`
+                : `const ${id}: ${formatConversions(varType, c)} = ${value};`
+}

package/src/node/scope.ts

@@ -3,24 +3,22 @@ import { conversion, node } from './node'
 import { getId } from './utils'
 import type { FnLayout, NodeProps, NodeProxy, X } from './types'
 
-let _scope: NodeProxy | null = null
-
-const scoped = (x: NodeProxy | null, fun = () => {}) => {
-        const prev = _scope
-        _scope = x
-        fun()
-        _scope = prev
-}
+let scope: NodeProxy | null = null
+let define: NodeProxy | null = null
 
 const addToScope = (x: NodeProxy) => {
-        if (!_scope) return // ignore
-        if (!_scope.props.children) _scope.props.children = []
-        _scope.props.children.push(x)
+        if (!scope) return
+        if (!scope.props.children) scope.props.children = []
+        scope.props.children.push(x)
+        if (x.type !== 'return' || !define) return
+        const { props } = define
+        if (!props.inferFrom) props.inferFrom = []
+        props.inferFrom.push(x)
 }
 
 export const toVar = (x: X, id?: string) => {
         if (!id) id = getId()
-        const y = node('variable', { id, inferFrom: x })
+        const y = node('variable', { id, inferFrom: [x] })
         const z = node('declare', null, x, y)
         addToScope(z)
         return y
@@ -32,6 +30,27 @@ export const assign = (x: X, y: X) => {
         return x
 }
 
+export const Return = (x: X) => {
+        const y = node('return', { inferFrom: [x] }, x)
+        addToScope(y)
+        return y
+}
+
+const scoped = (x: NodeProxy, fun: () => NodeProxy | void, y = define) => {
+        // cache to revert
+        const _scope = scope
+        const _define = define
+        // update
+        scope = x
+        define = y
+        if (_scope) x.props.parent = _scope
+        const z = fun()
+        if (z) Return(z)
+        // revert
+        scope = _scope
+        define = _define
+}
+
 export const If = (x: X, fun: () => void) => {
         const y = node('scope')
         scoped(y, fun)
@@ -55,7 +74,7 @@ export const If = (x: X, fun: () => void) => {
 
 export const Loop = (x: X, fun: (params: { i: NodeProxy }) => void) => {
         const y = node('scope')
-        scoped(y, () => fun({ i: node('variable', { id: 'i', inferFrom: int(0) }) }))
+        scoped(y, () => fun({ i: node('variable', { id: 'i', inferFrom: [int(0)] }) }))
         const ret = node('loop', null, x, y)
         addToScope(ret)
         return ret
@@ -88,20 +107,20 @@ export const Fn = (fun: (paramVars: NodeProxy[]) => NodeProxy) => {
         const ret = (...args: X[]) => {
                 const id = layout?.name || getId()
                 const x = node('scope')
-                let y: NodeProxy | undefined
                 const paramVars: NodeProxy[] = []
                 const paramDefs: NodeProps[] = []
                 if (layout?.inputs)
                         for (const input of layout.inputs) {
-                                paramDefs.push({ id: input.name, inferFrom: conversion(input.type) })
+                                paramDefs.push({ id: input.name, inferFrom: [conversion(input.type)] })
                         }
                 else
                         for (let i = 0; i < args.length; i++) {
-                                paramDefs.push({ id: `p${i}`, inferFrom: args[i] })
+                                paramDefs.push({ id: `p${i}`, inferFrom: [args[i]] })
                         }
                 for (const props of paramDefs) paramVars.push(node('variable', props))
-                scoped(x, () => (y = fun(paramVars)))
-                return node('define', { id, layout }, x, y, ...args)
+                const y = node('define', { id, layout }, x, ...args)
+                scoped(x, () => fun(paramVars), y)
+                return y
         }
         ret.setLayout = (newLayout: FnLayout) => {
                 layout = newLayout