--- a/src/zamikongine/shader.nim	Sat Jan 14 14:15:50 2023 +0700
+++ b/src/zamikongine/shader.nim	Sat Jan 14 23:34:50 2023 +0700
@@ -5,14 +5,16 @@
 
 
 import ./vulkan_helpers
+import ./glsl_helpers
 import ./vulkan
 import ./vertex
+import ./descriptor
 import ./math/vector
 
 type
   AllowedUniformType = SomeNumber|Vec
   UniformSlot *[T:AllowedUniformType] = object
-  ShaderProgram*[Uniforms] = object
+  ShaderProgram*[VertexType, Uniforms] = object
     entryPoint*: string
     programType*: VkShaderStageFlagBits
     shader*: VkPipelineShaderStageCreateInfo
@@ -58,7 +60,7 @@
     )
     i += 4
 
-proc initShaderProgram*[T](device: VkDevice, programType: static VkShaderStageFlagBits, shader: static string, entryPoint: static string="main"): ShaderProgram[T] =
+proc initShaderProgram*[VertexType, Uniforms](device: VkDevice, programType: static VkShaderStageFlagBits, shader: static string, entryPoint: static string="main"): ShaderProgram[VertexType, Uniforms] =
   result.entryPoint = entryPoint
   result.programType = programType
 
@@ -79,31 +81,63 @@
     pName: cstring(result.entryPoint), # entry point for shader
   )
 
-func generateVertexShaderCode*[VertexType](entryPoint, positionAttrName, colorAttrName: static string): string {.compileTime.} =
+func generateVertexShaderCode*[VertexType, Uniforms](
+  shaderBody: static string = "",
+  entryPoint: static string = "main",
+  glslVersion: static string = "450"
+): string {.compileTime.} =
   var lines: seq[string]
-  lines.add "#version 450"
-  # lines.add "layout(binding = 0) uniform UniformBufferObject { float dt; } ubo;"
-  lines.add generateGLSLDeclarations[VertexType]()
+  lines.add "#version " & glslVersion
+  lines.add "layout(row_major) uniform;"
+  lines.add generateGLSLUniformDeclarations[Uniforms]()
+  lines.add generateGLSLVertexDeclarations[VertexType]()
   lines.add "layout(location = 0) out vec3 fragColor;"
   lines.add "void " & entryPoint & "() {"
 
+  var hasPosition = 0
+  var hasColor = 0
   for name, value in VertexType().fieldPairs:
-    when typeof(value) is VertexAttribute and name == positionAttrName:
-      # lines.add "    vec2 tmp = " & name & " * ubo.dt;"
-      lines.add "    vec2 tmp = " & name & ";"
-      lines.add "    gl_Position = vec4(tmp, 0.0, 1.0);"
-    when typeof(value) is VertexAttribute and name == colorAttrName:
-      lines.add "    fragColor = " & name & ";"
+    when typeof(value) is PositionAttribute:
+      let glsltype = getGLSLType[getAttributeType(value)]()
+      lines.add &"    {glsltype} in_position = " & name & ";"
+      if getAttributeType(value) is Vec2:
+        lines.add "    vec4 out_position = vec4(in_position, 0.0, 1.0);"
+      elif getAttributeType(value) is Vec3:
+        lines.add "    vec4 out_position = vec4(in_position, 1.0);"
+      elif getAttributeType(value) is Vec4:
+        lines.add "    vec4 out_position = in_position;"
+      hasPosition += 1
+    when typeof(value) is ColorAttribute:
+      let glsltype = getGLSLType[getAttributeType(value)]()
+      lines.add &"    {glsltype} in_color = " & name & ";"
+      lines.add &"    {glsltype} out_color = in_color;";
+      hasColor += 1
+
+  lines.add shaderBody
+  lines.add "    gl_Position = out_position;"
+  lines.add "    fragColor = out_color;"
   lines.add "}"
+  if hasPosition != 1:
+    raise newException(Exception, fmt"VertexType needs to have exactly one attribute of type PositionAttribute (has {hasPosition})")
+  if hasColor != 1:
+    raise newException(Exception, fmt"VertexType needs to have exactly one attribute of type ColorAttribute (has {hasColor})")
   return lines.join("\n")
 
-func generateFragmentShaderCode*[VertexType](entryPoint: static string): string {.compileTime.} =
+func generateFragmentShaderCode*[VertexType](
+  shaderBody: static string = "",
+  entryPoint: static string = "main",
+  glslVersion: static string = "450"
+): string {.compileTime.} =
   var lines: seq[string]
-  lines.add "#version 450"
+  lines.add "#version " & glslVersion
+  lines.add "layout(row_major) uniform;"
   lines.add "layout(location = 0) in vec3 fragColor;"
   lines.add "layout(location = 0) out vec4 outColor;"
   lines.add "void " & entryPoint & "() {"
-  lines.add "    outColor = vec4(fragColor, 1.0);"
+  lines.add "    vec3 in_color = fragColor;"
+  lines.add "    vec3 out_color = in_color;"
+  lines.add shaderBody
+  lines.add "    outColor = vec4(out_color, 1.0);"
   lines.add "}"
 
   return lines.join("\n")