rlsl 0.1.0

data/lib/rlsl/prism/type_inference.rb

@@ -0,0 +1,289 @@
+# frozen_string_literal: true
+
+module RLSL
+  module Prism
+    class TypeInference
+      attr_reader :symbol_table
+
+      def initialize(uniforms = {}, custom_functions = {})
+        @symbol_table = {}
+        @uniforms = uniforms
+        @custom_functions = custom_functions
+
+        uniforms.each do |name, type|
+          @symbol_table[name.to_sym] = type
+        end
+      end
+
+      def register(name, type)
+        @symbol_table[name.to_sym] = type
+      end
+
+      def register_function(name, returns:)
+        @custom_functions[name.to_sym] = { returns: returns }
+      end
+
+      def lookup(name)
+        @symbol_table[name.to_sym]
+      end
+
+      def infer(node)
+        case node
+        when IR::Block
+          infer_block(node)
+        when IR::VarDecl
+          infer_var_decl(node)
+        when IR::VarRef
+          infer_var_ref(node)
+        when IR::Literal
+          infer_literal(node)
+        when IR::BoolLiteral
+          node.type = :bool
+          node
+        when IR::BinaryOp
+          infer_binary_op(node)
+        when IR::UnaryOp
+          infer_unary_op(node)
+        when IR::FuncCall
+          infer_func_call(node)
+        when IR::FieldAccess
+          infer_field_access(node)
+        when IR::Swizzle
+          infer_swizzle(node)
+        when IR::IfStatement
+          infer_if_statement(node)
+        when IR::Return
+          infer_return(node)
+        when IR::Assignment
+          infer_assignment(node)
+        when IR::ForLoop
+          infer_for_loop(node)
+        when IR::WhileLoop
+          infer_while_loop(node)
+        when IR::Parenthesized
+          infer_parenthesized(node)
+        when IR::FunctionDefinition
+          infer_function_definition(node)
+        when IR::ArrayLiteral
+          infer_array_literal(node)
+        when IR::ArrayIndex
+          infer_array_index(node)
+        when IR::GlobalDecl
+          infer_global_decl(node)
+        when IR::MultipleAssignment
+          infer_multiple_assignment(node)
+        else
+          node
+        end
+      end
+
+      private
+
+      def infer_block(node)
+        node.statements.each { |stmt| infer(stmt) }
+        node.type = node.statements.last&.type
+        node
+      end
+
+      def infer_var_decl(node)
+        infer(node.initializer) if node.initializer
+        node.type ||= node.initializer&.type
+        register(node.name, node.type) if node.type
+        node
+      end
+
+      def infer_var_ref(node)
+        node.type ||= lookup(node.name)
+        node
+      end
+
+      def infer_literal(node)
+        node
+      end
+
+      def infer_binary_op(node)
+        infer(node.left)
+        infer(node.right)
+
+        node.type = Builtins.binary_op_result_type(
+          node.operator,
+          node.left.type,
+          node.right.type
+        )
+        node
+      end
+
+      def infer_unary_op(node)
+        infer(node.operand)
+
+        case node.operator.to_s
+        when "-"
+          node.type = node.operand.type
+        when "!"
+          node.type = :bool
+        end
+        node
+      end
+
+      def infer_func_call(node)
+        node.args.each { |arg| infer(arg) }
+        infer(node.receiver) if node.receiver
+
+        sig = Builtins.function_signature(node.name)
+        if sig
+          arg_types = node.args.map(&:type)
+          node.type = Builtins.resolve_return_type(sig[:returns], arg_types)
+        elsif @custom_functions.key?(node.name.to_sym)
+          node.type = @custom_functions[node.name.to_sym][:returns]
+        else
+          node.type = node.receiver&.type
+        end
+        node
+      end
+
+      def infer_field_access(node)
+        infer(node.receiver)
+
+        if Builtins.single_component_field?(node.field)
+          node.type = :float
+        else
+          node.type = @uniforms[node.field.to_sym] || :float
+        end
+        node
+      end
+
+      def infer_swizzle(node)
+        infer(node.receiver)
+        node.type = Builtins.swizzle_type(node.components)
+        node
+      end
+
+      def infer_if_statement(node)
+        infer(node.condition)
+        infer(node.then_branch)
+        infer(node.else_branch) if node.else_branch
+
+        node.type = node.then_branch.type
+        node
+      end
+
+      def infer_return(node)
+        infer(node.expression) if node.expression
+        node.type = node.expression&.type
+        node
+      end
+
+      def infer_assignment(node)
+        infer(node.target)
+        infer(node.value)
+        node.type = node.value.type
+        node
+      end
+
+      def infer_for_loop(node)
+        register(node.variable, :int)
+        infer(node.range_start)
+        infer(node.range_end)
+        infer(node.body)
+        node.type = nil
+        node
+      end
+
+      def infer_while_loop(node)
+        infer(node.condition)
+        infer(node.body)
+        node.type = nil
+        node
+      end
+
+      def infer_function_definition(node)
+        node.param_types.each do |param_name, param_type|
+          register(param_name, param_type)
+        end
+
+        infer(node.body)
+
+        node.return_type ||= node.body&.type
+        node.type = node.return_type
+        node
+      end
+
+      def infer_parenthesized(node)
+        infer(node.expression)
+        node.type = node.expression.type
+        node
+      end
+
+      def infer_array_literal(node)
+        node.elements.each { |elem| infer(elem) }
+
+        element_type = node.elements.first&.type || :float
+        node.type = :"array_#{element_type}"
+        node
+      end
+
+      def infer_array_index(node)
+        infer(node.array)
+        infer(node.index)
+
+        array_type = node.array.type
+        if array_type.to_s.start_with?("array_")
+          node.type = array_type.to_s.sub("array_", "").to_sym
+        else
+          node.type = @symbol_table["#{node.array.name}_element_type".to_sym] || :float
+        end
+        node
+      end
+
+      def infer_global_decl(node)
+        infer(node.initializer) if node.initializer
+
+        if node.initializer.is_a?(IR::ArrayLiteral)
+          node.array_size ||= node.initializer.elements.length
+          first_elem = node.initializer.elements.first
+          node.element_type ||= first_elem&.type || :float
+          node.type = :"array_#{node.element_type}"
+        else
+          node.type ||= node.initializer&.type
+        end
+
+        register(node.name, node.type) if node.type
+
+        if node.element_type
+          register("#{node.name}_element_type".to_sym, node.element_type)
+        end
+
+        node
+      end
+
+      def infer_multiple_assignment(node)
+        infer(node.value)
+
+        value_type = node.value.type
+        if value_type.is_a?(IR::TupleType)
+          node.targets.each_with_index do |target, i|
+            target.type = value_type.types[i]
+            register(target.name, target.type)
+          end
+        elsif value_type.to_s.start_with?("array_")
+          elem_type = value_type.to_s.sub("array_", "").to_sym
+          node.targets.each do |target|
+            target.type = elem_type
+            register(target.name, target.type)
+          end
+        elsif @custom_functions.key?(node.value.name) && node.value.is_a?(IR::FuncCall)
+          func_info = @custom_functions[node.value.name]
+          if func_info[:returns].is_a?(Array)
+            node.targets.each_with_index do |target, i|
+              target.type = func_info[:returns][i]
+              register(target.name, target.type)
+            end
+          end
+        end
+
+        node.type = nil
+        node
+      end
+    end
+  end
+end

data/lib/rlsl/shader_builder.rb

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "digest"
+require "rbconfig"
+
+module RLSL
+  class ShaderBuilder
+    attr_reader :name
+
+    def initialize(name)
+      @name = name.to_s
+      @uniforms = {}
+      @fragment_mode = :c
+      @helpers_mode = :c
+      @custom_functions = {}
+    end
+
+    def uniforms(&block)
+      if block_given?
+        ctx = UniformContext.new
+        ctx.instance_eval(&block)
+        @uniforms = ctx.uniforms
+      else
+        @uniforms
+      end
+    end
+
+    def helpers(mode = :ruby, &block)
+      @helpers_block = block
+      @helpers_mode = mode
+    end
+
+    def functions(&block)
+      ctx = FunctionContext.new
+      ctx.instance_eval(&block)
+      @custom_functions = ctx.functions
+    end
+
+    def fragment(&block)
+      @fragment_block = block
+      @fragment_mode = block.arity > 0 ? :ruby : :c
+    end
+
+    def ruby_mode?
+      @fragment_mode == :ruby
+    end
+
+    def compile_and_load
+      c_code = generate_c_code
+      code_hash = Digest::MD5.hexdigest(c_code)[0..7]
+      ext_name = "#{@name}_#{code_hash}"
+      ext_dir = File.join(RLSL.cache_dir, ext_name)
+      ext_file = File.join(ext_dir, "#{@name}.#{RbConfig::CONFIG['DLEXT']}")
+
+      unless File.exist?(ext_file)
+        compile_extension(@name, ext_dir, c_code)
+      end
+
+      require ext_file
+      CompiledShader.new(@name, ext_name, @uniforms.keys)
+    end
+
+    def build_metal_shader
+      if ruby_mode?
+        fragment_code = transpile_fragment(:msl)
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+      else
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+        fragment_code = @fragment_block ? @fragment_block.call : ""
+      end
+
+      translator = MSL::Translator.new(@uniforms, helpers_code, fragment_code)
+      msl_source = translator.translate
+
+      MSL::Shader.new(@name, @uniforms, msl_source)
+    end
+
+    def build_wgsl_shader
+      if ruby_mode?
+        fragment_code = transpile_fragment(:wgsl)
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+      else
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+        fragment_code = @fragment_block ? @fragment_block.call : ""
+      end
+
+      translator = WGSL::Translator.new(@uniforms, helpers_code, fragment_code)
+      translator.translate
+    end
+
+    def build_glsl_shader(version: "450")
+      if ruby_mode?
+        fragment_code = transpile_fragment(:glsl)
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+      else
+        helpers_code = @helpers_block ? @helpers_block.call : ""
+        fragment_code = @fragment_block ? @fragment_block.call : ""
+      end
+
+      translator = GLSL::Translator.new(@uniforms, helpers_code, fragment_code, version: version)
+      translator.translate
+    end
+
+    def transpile_fragment(target)
+      return "" unless @fragment_block
+
+      transpiler = Prism::Transpiler.new(@uniforms, @custom_functions)
+      transpiler.transpile(@fragment_block, target)
+    end
+
+    def transpile_helpers(target)
+      return "" unless @helpers_block
+
+      transpiler = Prism::Transpiler.new(@uniforms, @custom_functions)
+      transpiler.transpile_helpers(@helpers_block, target, @custom_functions)
+    end
+
+    def helpers_ruby_mode?
+      @helpers_mode == :ruby
+    end
+
+    private
+
+    def generate_c_code
+      if helpers_ruby_mode?
+        helpers_code = transpile_helpers(:c)
+        helpers_block = -> { helpers_code }
+      else
+        helpers_block = @helpers_block
+      end
+
+      if ruby_mode?
+        fragment_code = transpile_fragment(:c)
+        fragment_block = -> { fragment_code }
+      else
+        fragment_block = @fragment_block
+      end
+
+      codegen = CodeGenerator.new(@name, @uniforms, helpers_block, fragment_block)
+      codegen.generate
+    end
+
+    def compile_extension(ext_name, ext_dir, c_code)
+      FileUtils.mkdir_p(ext_dir)
+
+      File.write(File.join(ext_dir, "#{ext_name}.c"), c_code)
+
+      extconf = <<~RUBY
+        require "mkmf"
+        $CFLAGS << " -O3 -ffast-math"
+        if RUBY_PLATFORM =~ /darwin/
+          $CFLAGS << " -fblocks"
+        end
+        create_makefile("#{ext_name}")
+      RUBY
+      File.write(File.join(ext_dir, "extconf.rb"), extconf)
+
+      Dir.chdir(ext_dir) do
+        system("#{RbConfig.ruby} extconf.rb > /dev/null 2>&1") or raise "extconf failed for #{ext_name}"
+        system("/usr/bin/make > /dev/null 2>&1") or raise "make failed for #{ext_name}"
+      end
+    end
+  end
+end

data/lib/rlsl/types.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module RLSL
+  # C type definitions for the shader system
+  C_TYPES = <<~C
+    typedef struct { float x, y; } vec2;
+    typedef struct { float x, y, z; } vec3;
+    typedef struct { float x, y, z, w; } vec4;
+    typedef struct { float m[4]; } mat2;
+    typedef struct { float m[9]; } mat3;
+    typedef struct { float m[16]; } mat4;
+    typedef struct { void* data; int width; int height; } sampler2D;
+
+    #define PI 3.14159265f
+    #define TAU 6.28318530f
+  C
+
+  # Uniform type symbols
+  UNIFORM_TYPES = %i[float vec2 vec3 vec4 int bool mat2 mat3 mat4 sampler2D].freeze
+
+  # Type mapping helper
+  module TypeMapping
+    C_UNIFORM_TYPES = {
+      float: "float",
+      int: "int",
+      bool: "int",
+      vec2: "vec2",
+      vec3: "vec3",
+      vec4: "vec4",
+      mat2: "mat2",
+      mat3: "mat3",
+      mat4: "mat4",
+      sampler2D: "sampler2D"
+    }.freeze
+  end
+end

data/lib/rlsl/uniform_context.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module RLSL
+  # DSL context for defining uniform variables
+  class UniformContext
+    attr_reader :uniforms
+
+    def initialize
+      @uniforms = {}
+    end
+
+    def float(name)
+      @uniforms[name] = :float
+    end
+
+    def vec2(name)
+      @uniforms[name] = :vec2
+    end
+
+    def vec3(name)
+      @uniforms[name] = :vec3
+    end
+
+    def vec4(name)
+      @uniforms[name] = :vec4
+    end
+  end
+end

data/lib/rlsl/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module RLSL
+  VERSION = "0.1.0"
+end

data/lib/rlsl/wgsl/translator.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+module RLSL
+  module WGSL
+    class Translator < BaseTranslator
+      TYPE_MAP = {
+        "vec2" => "vec2<f32>",
+        "vec3" => "vec3<f32>",
+        "vec4" => "vec4<f32>"
+      }.freeze
+
+      FUNC_REPLACEMENTS = BaseTranslator.common_func_replacements(
+        target_vec2: "vec2<f32>",
+        target_vec3: "vec3<f32>",
+        target_vec4: "vec4<f32>"
+      ).freeze
+
+      protected
+
+      def translate_code(c_code)
+        result = super(c_code)
+        return result if result.empty?
+
+        result.gsub!(/\bstatic\s+/, "")
+        result.gsub!(/\binline\s+/, "")
+        result.gsub!(/\bfloat\b/, "f32")
+        result
+      end
+
+      def generate_shader(helpers, fragment)
+        <<~WGSL
+          // WGSL Shader generated by RLSL
+
+          struct Uniforms {
+              #{generate_uniform_struct}
+          };
+
+          @group(0) @binding(0) var<uniform> u: Uniforms;
+          @group(0) @binding(1) var output_texture: texture_storage_2d<rgba8unorm, write>;
+
+          #{helpers}
+
+          fn shader_fragment(frag_coord: vec2<f32>, resolution: vec2<f32>) -> vec3<f32> {
+              let uv = frag_coord / resolution.y;
+              #{fragment}
+          }
+
+          @compute @workgroup_size(8, 8)
+          fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+              let resolution = vec2<f32>(f32(textureDimensions(output_texture).x),
+                                         f32(textureDimensions(output_texture).y));
+              if (gid.x >= u32(resolution.x) || gid.y >= u32(resolution.y)) {
+                  return;
+              }
+
+              let frag_coord = vec2<f32>(f32(gid.x), resolution.y - 1.0 - f32(gid.y));
+              let color = shader_fragment(frag_coord, resolution);
+
+              textureStore(output_texture, vec2<i32>(i32(gid.x), i32(gid.y)),
+                           vec4<f32>(clamp(color, vec3<f32>(0.0), vec3<f32>(1.0)), 1.0));
+          }
+        WGSL
+      end
+
+      private
+
+      def generate_uniform_struct
+        fields = ["resolution: vec2<f32>,"]
+        @uniforms.each do |name, type|
+          wgsl_type = uniform_type_to_target(type)
+          fields << "#{name}: #{wgsl_type},"
+        end
+        fields.join("\n    ")
+      end
+
+      def target_float_type
+        "f32"
+      end
+
+      def target_vec2_type
+        "vec2<f32>"
+      end
+
+      def target_vec3_type
+        "vec3<f32>"
+      end
+
+      def target_vec4_type
+        "vec4<f32>"
+      end
+    end
+  end
+end

data/lib/rlsl.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require "fileutils"
+
+require_relative "rlsl/version"
+require_relative "rlsl/types"
+require_relative "rlsl/uniform_context"
+require_relative "rlsl/function_context"
+require_relative "rlsl/code_generator"
+require_relative "rlsl/compiled_shader"
+require_relative "rlsl/base_translator"
+require_relative "rlsl/msl/translator"
+require_relative "rlsl/msl/shader"
+require_relative "rlsl/wgsl/translator"
+require_relative "rlsl/glsl/translator"
+require_relative "rlsl/prism/transpiler"
+require_relative "rlsl/shader_builder"
+
+module RLSL
+  CACHE_DIR = File.expand_path("~/.cache/rlsl/compiled")
+
+  class << self
+    def define(name, &block)
+      builder = ShaderBuilder.new(name)
+      builder.instance_eval(&block)
+      builder.compile_and_load
+    end
+
+    def define_metal(name, &block)
+      builder = ShaderBuilder.new(name)
+      builder.instance_eval(&block)
+      builder.build_metal_shader
+    end
+
+    def to_wgsl(name, &block)
+      builder = ShaderBuilder.new(name)
+      builder.instance_eval(&block)
+      builder.build_wgsl_shader
+    end
+
+    def to_glsl(name, version: "450", &block)
+      builder = ShaderBuilder.new(name)
+      builder.instance_eval(&block)
+      builder.build_glsl_shader(version: version)
+    end
+
+    def cache_dir
+      @cache_dir ||= begin
+        FileUtils.mkdir_p(CACHE_DIR)
+        CACHE_DIR
+      end
+    end
+  end
+
+  module CompiledShaders
+  end
+end