Merge pull request #1 from arhik/main

`decl` expressions and `type` expressions

Author: arhik

src/WGPUTranspiler.jl

@@ -12,8 +12,8 @@ include("codegen/rangeBlock.jl")
 include("codegen/conditionBlock.jl")
 include("codegen/funcBlock.jl")
 include("codegen/builtin.jl")
-include("codegen/computeBlock.jl")
 include("codegen/expr.jl")
+include("codegen/computeBlock.jl")
 include("codegen/infer.jl")
 include("codegen/resolve.jl")
 include("codegen/transpile.jl")

src/codegen/assignment.jl

@@ -40,7 +40,7 @@ end
 symbol(assign::AssignmentExpr) = symbol(assign.lhs)
 
 function assignExpr(scope, lhs, rhs)
-	lhsVar = LHS(inferVariable(scope, lhs), false)
+	lhsVar = LHS(inferExpr(scope, lhs), false)
 	inferScope!(scope, lhsVar)
 	rhsExpr = RHS(inferExpr(scope, rhs))
 	statement = AssignmentExpr(lhsVar, rhsExpr, scope)

src/codegen/computeBlock.jl

@@ -23,7 +23,7 @@ end
 
 struct ComputeBlock <: JLBlock
 	fname::WGPUVariable
-	fargs::Vector{WGPUVariable}
+	fargs::Vector{DeclExpr}
 	Targs::Vector{WGPUVariable}
 	fbody::Vector{JLExpr}
 	scope::Union{Nothing, Scope}
@@ -35,7 +35,7 @@ function computeBlock(scope, islaunch, wgSize, wgCount, fname, fargs)
 	@capture(fexpr, function fname_(fargs__) where Targs__ fbody__ end)
 	childScope = Scope([Targs...], [:ceil], 0, scope, quote end)
 	fn = inferExpr(childScope, fname)
-	fa = map(x -> inferVariable(childScope, x), fargs)
+	fa = map(x -> inferExpr(childScope, x), fargs)
 	fb = map(x -> inferExpr(childScope, x), fbody)
 	return ComputeBlock(fn, fa, WGPUVariable[], fb, childScope)
 end

src/codegen/expr.jl

@@ -102,5 +102,22 @@ end
 symbol(access::AccessExpr) = symbol(access.sym)
 
 function inferScope!(scope::Scope, jlexpr::AccessExpr)
+	#inferScope!(scope, jlexpr.sym)
+	#inferScope!(scope, jlexpr.field)
+end
+
+struct TypeExpr <: JLExpr
+	sym::WGPUVariable
+	types::Vector{WGPUVariable}
+end
+
+symbol(tExpr::TypeExpr) = (symbol(tExpr.sym), map(x -> symbol(x), tExpr.types)...)
 
+struct DeclExpr <: JLExpr
+	sym::WGPUVariable
+	dataType::Union{DataType, TypeExpr}
 end
+
+symbol(decl::DeclExpr) = symbol(decl.sym)
+
+

src/codegen/infer.jl

@@ -29,10 +29,14 @@ function inferExpr(scope::Scope, expr::Expr)
 		return binaryOp(scope, :-=, a, b)
 	elseif @capture(expr, f_(args__))
 		return callExpr(scope, f, args)
+	elseif @capture(expr, a_::b_)
+		return declExpr(scope, a, b)
 	elseif @capture(expr, a_[b_])
 		return indexExpr(scope, a, b)
 	elseif @capture(expr, a_.b_)
 		return accessExpr(scope, a, b)
+	elseif @capture(expr, a_{b__})
+		return typeExpr(scope, a, b)
 	elseif @capture(expr, for idx_ in range_ block__ end)
 		return rangeBlock(scope, idx, range, block)
 	elseif @capture(expr, if cond_ block__ end)
@@ -52,15 +56,37 @@ function inferExpr(scope::Scope, expr::Expr)
 	end
 end
 
+function declExpr(scope, a::Val{:hello}) 
+	@error "Not implemented yet"
+end
+
+function declExpr(scope, a::Symbol, b::Symbol)
+	aExpr = inferExpr(scope, a)
+	inferScope!(scope, aExpr)
+	bExpr = eval(b)
+	return DeclExpr(aExpr, bExpr)
+end
+
+function typeExpr(scope, a::Symbol, b::Vector{Any})
+	aExpr = inferExpr(scope, a)
+	bExpr = map(x -> inferExpr(scope, x), b)
+	return TypeExpr(aExpr, bExpr)
+end
+
+function declExpr(scope, a::Symbol, b::Expr)
+	bExpr = inferExpr(scope, b)
+	aExpr = inferExpr(scope, a)
+	return DeclExpr(aExpr, bExpr)
+end
+
+
 function inferVariable(scope, expr::Expr)
 	if @capture(expr, a_::b_{t__})
 		push!(scope.locals, a)
 		return WGPUVariable(a, eval(b), Generic, nothing, ) # TODO t is ignored
 	elseif @capture(expr, a_::b_)
 		push!(scope.locals, a)
 		return WGPUVariable(a, eval(b), Generic, nothing, )
-	elseif @capture(expr, a_[b_])
-		return indexExpr(scope, a, b)
 	else
 		error("This expression $expr type is not captured yet")
 	end

src/codegen/test.jl

@@ -89,7 +89,7 @@ end
 a = WgpuArray(rand(Float32, 4, 4));
 b = WgpuArray(rand(Float32, 4, 4));
 
-scope = Scope([], [], 0, nothing, quote end)
+scope = Scope([:out, :x], [], 0, nothing, quote end)
 inferredExpr = inferExpr(
 	scope, 
 	:(@wgpukernel launch=true workgroupSize=(4, 4) workgroupCount=(1, 1) $cast_kernel($a, $b))

src/codegen/transpile.jl

@@ -2,9 +2,9 @@ export transpile
 
 transpile(scope::Scope, s::Scalar) = s.element
 transpile(scope::Scope, var::WGPUVariable) = var.dataType == Any ? :($(var.sym)) : :($(var.sym)::$(var.dataType))
-transpile(scope::Scope, lhs::LHS) = transpile(scope, lhs.variable, Val(lhs.mutable))
-transpile(scope::Scope, var::WGPUVariable, ::Val{false}) = :(@var $(transpile(scope, var)))
-transpile(scope::Scope, var::WGPUVariable, ::Val{true}) = :($(var.sym))
+transpile(scope::Scope, lhs::LHS) = transpile(scope, lhs.variable)
+transpile(scope::Scope, var::WGPUVariable, ::Val{true}) = :(@var $(transpile(scope, var)))
+transpile(scope::Scope, var::WGPUVariable, ::Val{false}) = :($(var.sym))
 transpile(scope::Scope, rhs::RHS) = transpile(scope, rhs.rhsExpr)
 transpile(scope::Scope, binOp::BinaryOp) = transpile(scope, binOp, Val(binOp.op))
 
@@ -46,6 +46,14 @@ function transpile(scope::Scope, acsExpr::AccessExpr)
 	return Expr(:., transpile(scope, acsExpr.sym), QuoteNode(transpile(scope, acsExpr.field)))
 end
 
+transpile(scope::Scope, declExpr::DeclExpr) = Expr(:(::), map(x -> transpile(scope, x), (declExpr.sym, declExpr.dataType))...)
+transpile(scope::Scope, ::Type{T}) where T = :($T)
+
+transpile(scope::Scope, typeExpr::TypeExpr) = Expr(
+	:curly, transpile(scope, typeExpr.sym), 
+	map(x -> transpile(scope, x), typeExpr.types)...
+)
+
 function transpile(scope::Scope, rblock::RangeBlock)
 	(start, step, stop) = map(x -> transpile(scope, x), (rblock.start, rblock.step, rblock.stop))
 	range = :($start:$step:$stop)