mt-lang 0.2.0

data/lib/milk_tea/core/control_flow/constant_propagation.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require_relative "../compile_time/const_eval"
+
+module MilkTea
+  module ControlFlow
+    class ConstantPropagation
+      ConstVal = Data.define(:value)
+
+      UNDEF = :undef
+      NAC   = :nac
+
+      Result = Data.define(:in_states, :out_states) do
+        def constant_at(node_id, key)
+          entry = in_states[node_id]
+          return nil unless entry
+
+          v = entry[key]
+          v.is_a?(ConstVal) ? v.value : nil
+        end
+      end
+
+      def self.solve(graph, binding_resolution: nil, strict_binding_ids: false)
+        result = Dataflow.solve(
+          graph,
+          direction: :forward,
+          initial: -> { {} },
+          join: lambda do |states|
+            return {} if states.empty?
+
+            states.reduce do |acc, state|
+              keys = acc.keys | state.keys
+              keys.each_with_object({}) do |k, merged|
+                a = acc[k]   || UNDEF
+                b = state[k] || UNDEF
+                merged[k] = join_lattice(a, b)
+              end
+            end
+          end,
+          transfer: lambda do |node, in_state|
+            out = in_state.dup
+            node.writes_info.each do |write|
+              key = write[:binding_key]
+              val = eval_const(node.statement, write, in_state, binding_resolution:, strict_binding_ids:)
+              out[key] = val
+            end
+            out
+          end,
+          boundary_in: { graph.entry_id => {} }
+        )
+
+        Result.new(in_states: result.in_states, out_states: result.out_states)
+      end
+
+      def self.join_lattice(a, b)
+        return b if a == UNDEF
+        return a if b == UNDEF
+        return NAC if a == NAC || b == NAC
+
+        a == b ? a : NAC
+      end
+      private_class_method :join_lattice
+
+      def self.eval_const(statement, _write, in_state, binding_resolution:, strict_binding_ids:)
+        case statement
+        when AST::LocalDecl
+          eval_expr_const(statement.value, in_state, binding_resolution:, strict_binding_ids:)
+        when AST::Assignment
+          eval_assignment_const(statement, in_state, binding_resolution:, strict_binding_ids:)
+        else
+          NAC
+        end
+      end
+      private_class_method :eval_const
+
+      def self.eval_assignment_const(statement, in_state, binding_resolution:, strict_binding_ids:)
+        return NAC unless statement.target.is_a?(AST::Identifier)
+
+        rhs = eval_expr_const(statement.value, in_state, binding_resolution:, strict_binding_ids:)
+        return NAC unless rhs.is_a?(ConstVal)
+
+        return rhs if statement.operator == "="
+
+        lhs = eval_expr_const(statement.target, in_state, binding_resolution:, strict_binding_ids:)
+        return NAC unless lhs.is_a?(ConstVal)
+
+        begin
+          v = case statement.operator
+              when "+="  then lhs.value + rhs.value
+              when "-="  then lhs.value - rhs.value
+              when "*="  then lhs.value * rhs.value
+              when "/="  then rhs.value.zero? ? (return NAC) : lhs.value / rhs.value
+              when "%="  then rhs.value.zero? ? (return NAC) : lhs.value % rhs.value
+              when "&="  then lhs.value & rhs.value
+              when "|="  then lhs.value | rhs.value
+              when "^="  then lhs.value ^ rhs.value
+              when "<<=" then lhs.value << rhs.value
+              when ">>=" then lhs.value >> rhs.value
+              else return NAC
+              end
+          ConstVal.new(v)
+        rescue ZeroDivisionError, RangeError
+          NAC
+        end
+      end
+      private_class_method :eval_assignment_const
+
+      def self.eval_expr_const(expr, state, binding_resolution:, strict_binding_ids:)
+        return NAC if expr.nil?
+
+        value = ConstEval.evaluate(
+          expr,
+          resolve_identifier: lambda do |identifier_expression|
+            key = identifier_key(identifier_expression, binding_resolution:, strict_binding_ids:)
+            next unless key
+
+            state_value = state[key]
+            state_value.is_a?(ConstVal) ? state_value.value : nil
+          end,
+          resolve_member_access: nil,
+        )
+        value.nil? ? NAC : ConstVal.new(value)
+      end
+      private_class_method :eval_expr_const
+
+      def self.constant_value_of(expr, in_state, binding_resolution: nil, strict_binding_ids: false)
+        result = send(:eval_expr_const, expr, in_state, binding_resolution:, strict_binding_ids:)
+        result.is_a?(ConstVal) ? result.value : nil
+      end
+
+      def self.identifier_key(identifier_expression, binding_resolution:, strict_binding_ids:)
+        if binding_resolution && (id = binding_resolution.identifier_binding_ids[identifier_expression.object_id])
+          return id
+        end
+
+        return nil if strict_binding_ids
+
+        identifier_expression.name
+      end
+      private_class_method :identifier_key
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/dataflow.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class Dataflow
+      Result = Data.define(:in_states, :out_states)
+
+      def self.solve(graph, direction:, initial:, join:, transfer: nil, edge_transfer: nil, boundary_in: {}, boundary_out: {})
+        raise ArgumentError, "direction must be :forward or :backward" unless %i[forward backward].include?(direction)
+        raise ArgumentError, "provide either transfer: or edge_transfer:" if transfer.nil? && edge_transfer.nil?
+
+        in_states  = {}
+        out_states = {}
+        edge_out   = {}  # [from_id, to_id] => state  (only when edge_transfer given)
+        all_ids = graph.ids
+        all_ids.each do |id|
+          in_states[id]  = initial.call
+          out_states[id] = initial.call
+        end
+
+        changed = true
+        while changed
+          changed = false
+          iteration_ids = direction == :forward ? graph.rpo_ids : all_ids.reverse
+          iteration_ids.each do |id|
+            node = graph.nodes[id]
+
+            if direction == :forward
+              incoming =
+                if boundary_in.key?(id)
+                  boundary_in[id]
+                elsif edge_transfer
+                  join.call(node.preds.map { |pred| edge_out[[pred, id]] || initial.call })
+                else
+                  join.call(node.preds.map { |pred| out_states[pred] })
+                end
+
+              if edge_transfer
+                node.succs.each do |succ|
+                  label     = graph.edge_label(id, succ)
+                  new_edge  = edge_transfer.call(node, incoming, succ, label)
+                  old_edge  = edge_out[[id, succ]]
+                  if old_edge != new_edge
+                    edge_out[[id, succ]] = new_edge
+                    changed = true
+                  end
+                end
+                new_out = node.succs.empty? ? initial.call : join.call(node.succs.map { |s| edge_out[[id, s]] || initial.call })
+              else
+                new_out = boundary_out.fetch(id) { transfer.call(node, incoming) }
+              end
+
+              if in_states[id] != incoming || out_states[id] != new_out
+                in_states[id]  = incoming
+                out_states[id] = new_out
+                changed = true
+              end
+            else
+              outgoing = boundary_out.fetch(id) { join.call(node.succs.map { |succ| in_states[succ] }) }
+              incoming = boundary_in.fetch(id)  { transfer.call(node, outgoing) }
+              if in_states[id] != incoming || out_states[id] != outgoing
+                in_states[id]  = incoming
+                out_states[id] = outgoing
+                changed = true
+              end
+            end
+          end
+        end
+
+        Result.new(in_states:, out_states:)
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/definite_assignment.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class DefiniteAssignment
+      ReadBeforeAssignment = Data.define(:node_id, :binding_key, :line, :column, :length)
+      Result = Data.define(:definitely_assigned_in, :definitely_assigned_out, :read_before_assignment)
+
+      def self.solve(graph, initially_assigned: Set.new)
+        initially_assigned = initially_assigned.dup
+        universe = initial_universe(graph, initially_assigned)
+
+        result = Dataflow.solve(
+          graph,
+          direction: :forward,
+          initial: -> { universe.dup },
+          join: lambda do |states|
+            if states.empty?
+              universe.dup
+            else
+              states.reduce(universe.dup) { |acc, state| acc & state }
+            end
+          end,
+          transfer: lambda do |node, in_state|
+            in_state | node.writes
+          end,
+          boundary_in: { graph.entry_id => initially_assigned.dup }
+        )
+
+        read_before_assignment = []
+        graph.each_node do |node|
+          in_state = result.in_states[node.id]
+          if node.reads_info.empty?
+            node.reads.each do |binding_key|
+              next if in_state.include?(binding_key)
+
+              read_before_assignment << ReadBeforeAssignment.new(
+                node_id: node.id,
+                binding_key:,
+                line: node.line,
+                column: nil,
+                length: nil,
+              )
+            end
+            next
+          end
+
+          node.reads_info.each do |read_site|
+            next if in_state.include?(read_site.binding_key)
+
+            read_before_assignment << ReadBeforeAssignment.new(
+              node_id: node.id,
+              binding_key: read_site.binding_key,
+              line: read_site.line || node.line,
+              column: read_site.column,
+              length: read_site.length,
+            )
+          end
+        end
+
+        Result.new(
+          definitely_assigned_in: result.in_states,
+          definitely_assigned_out: result.out_states,
+          read_before_assignment:
+        )
+      end
+
+      def self.initial_universe(graph, initially_assigned)
+        universe = initially_assigned.dup
+        graph.each_node do |node|
+          universe.merge(node.writes)
+          universe.merge(node.reads)
+        end
+        universe
+      end
+      private_class_method :initial_universe
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/graph.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class Graph
+      attr_accessor :entry_id, :exit_id
+      attr_reader :nodes
+
+      def initialize
+        @nodes = {}
+        @next_id = 1
+        @entry_id = nil
+        @exit_id = nil
+        @edge_labels = {}      # Hash[(from_id, to_id) → :true_branch | :false_branch]
+        @edge_refinements = {} # Hash[(from_id, to_id) → Hash[binding_key → :non_null | :null]]
+      end
+
+      def add_node(kind:, statement: nil, line: nil, reads: Set.new, reads_info: [], writes: Set.new, writes_info: [])
+        id = @next_id
+        @next_id += 1
+        @nodes[id] = Node.new(
+          id:,
+          kind:,
+          statement:,
+          line:,
+          reads: reads.dup,
+          reads_info: reads_info.map(&:dup),
+          writes: writes.dup,
+          writes_info: writes_info.map(&:dup),
+          succs: [],
+          preds: []
+        )
+        id
+      end
+
+      def add_edge(from, to, label: nil)
+        return unless @nodes[from] && @nodes[to]
+
+        @nodes[from].succs << to unless @nodes[from].succs.include?(to)
+        @nodes[to].preds << from unless @nodes[to].preds.include?(from)
+        @edge_labels[[from, to]] = label if label
+      end
+
+      def edge_label(from, to)
+        @edge_labels[[from, to]]
+      end
+
+      def edge_refinement(from, to)
+        @edge_refinements[[from, to]]
+      end
+
+      def set_edge_refinement(from, to, refinement)
+        @edge_refinements[[from, to]] = refinement
+      end
+
+      def each_node
+        return enum_for(:each_node) unless block_given?
+
+        @nodes.each_value { |node| yield node }
+      end
+
+      def ids
+        @nodes.keys
+      end
+
+      def rpo_ids
+        visited = {}
+        order = []
+
+        dfs = lambda do |id|
+          return if visited[id]
+
+          visited[id] = true
+          node = @nodes[id]
+          node&.succs&.each { |succ| dfs.call(succ) }
+          order.unshift(id)
+        end
+
+        dfs.call(@entry_id) if @entry_id
+        order
+      end
+
+      def read_bindings
+        keys = Set.new
+        each_node { |node| keys.merge(node.reads) }
+        keys
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/liveness.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class Liveness
+      Result = Data.define(:live_in, :live_out)
+
+      def self.solve(graph)
+        result = Dataflow.solve(
+          graph,
+          direction: :backward,
+          initial: -> { Set.new },
+          join: lambda do |states|
+            states.reduce(Set.new) { |acc, state| acc | state }
+          end,
+          transfer: lambda do |node, live_out|
+            node.reads | (live_out - node.writes)
+          end,
+        )
+        Result.new(live_in: result.in_states, live_out: result.out_states)
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/nullability_flow.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class NullabilityFlow
+      Result = Data.define(:in_states, :out_states, :stmt_to_node_id) do
+        def nonnull_before(stmt)
+          node_id = stmt_to_node_id[stmt.object_id]
+          return Set.new unless node_id
+
+          in_states[node_id] || Set.new
+        end
+      end
+
+      def self.solve(graph)
+        stmt_to_node_id = {}
+        graph.each_node { |n| stmt_to_node_id[n.statement.object_id] = n.id if n.statement }
+
+        result = Dataflow.solve(
+          graph,
+          direction: :forward,
+          initial: -> { Set.new },
+          join: lambda do |states|
+            return Set.new if states.empty?
+
+            states.reduce { |acc, s| acc & s }
+          end,
+          edge_transfer: lambda do |node, in_state, succ_id, _edge_label|
+            state = in_state - node.writes
+            refs  = graph.edge_refinement(node.id, succ_id) || {}
+            refs.each do |key, ref|
+              ref == :non_null ? (state = state | Set[key]) : (state = state - Set[key])
+            end
+            state
+          end,
+          boundary_in: { graph.entry_id => Set.new }
+        )
+
+        Result.new(in_states: result.in_states, out_states: result.out_states, stmt_to_node_id:)
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/reachability.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class Reachability
+      Result = Data.define(:reachable_ids)
+
+      def self.solve(graph)
+        reachable = Set.new
+        queue = [graph.entry_id]
+        until queue.empty?
+          id = queue.shift
+          next if reachable.include?(id)
+
+          reachable << id
+          (graph.nodes[id]&.succs || []).each { |succ| queue << succ }
+        end
+        Result.new(reachable_ids: reachable)
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow/termination.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module ControlFlow
+    class Termination
+      def self.block_always_terminates?(statements, **builder_options)
+        return false if statements.nil? || statements.empty?
+
+        graph = Builder.new(**builder_options).build(statements)
+        reachability = Reachability.solve(graph)
+        !reachability.reachable_ids.include?(graph.exit_id)
+      end
+
+      def self.loop_body_always_exits?(statements, **builder_options)
+        return false if statements.nil? || statements.empty?
+
+        graph = Builder.new(**builder_options).build_loop_body(statements)
+        reachability = Reachability.solve(graph)
+        !reachability.reachable_ids.include?(graph.exit_id)
+      end
+
+      def self.block_always_terminates_in_loop?(statements, **builder_options)
+        return false if statements.nil? || statements.empty?
+
+        graph = Builder.new(**builder_options).build_loop_branch(statements)
+        reachability = Reachability.solve(graph)
+        !reachability.reachable_ids.include?(graph.exit_id)
+      end
+    end
+  end
+end

data/lib/milk_tea/core/control_flow.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+require "set"
+
+require_relative "control_flow/graph"
+require_relative "control_flow/builder"
+require_relative "control_flow/dataflow"
+require_relative "control_flow/liveness"
+require_relative "control_flow/definite_assignment"
+require_relative "control_flow/reachability"
+require_relative "control_flow/termination"
+require_relative "control_flow/nullability_flow"
+require_relative "control_flow/constant_propagation"
+
+module MilkTea
+  module ControlFlow
+    BindingResolution = Data.define(:identifier_binding_ids, :declaration_binding_ids, :mutating_argument_identifier_ids)
+    ReadSite = Data.define(:binding_key, :line, :column, :length)
+
+    Node = Struct.new(
+      :id,
+      :kind,
+      :statement,
+      :line,
+      :reads,
+      :reads_info,
+      :writes,
+      :writes_info,
+      :succs,
+      :preds,
+      keyword_init: true
+    )
+  end
+end

data/lib/milk_tea/core/cst.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module MilkTea
+  module CST
+    SourceFile = Data.define(:source, :tokens, :trivia) do
+      def reconstruct
+        reconstruct_from_tokens
+      end
+
+      def reconstruct_normalized
+        reconstruct_from_tokens(normalize_spaces: true)
+      end
+
+      def reconstruct_from_tokens(normalize_spaces: false)
+        return "" if tokens.empty?
+
+        tokens.each_with_object(+"") do |token, result|
+          append_trivia(result, token.leading_trivia, normalize_spaces:, leading: true)
+          result << segment_text(token.start_offset, token.end_offset, token.lexeme) unless token.eof?
+          append_trivia(result, token.trailing_trivia, normalize_spaces:, leading: false)
+        end
+      end
+
+      def append_trivia(result, entries, normalize_spaces:, leading:)
+        entries.each do |entry|
+          if normalize_spaces && entry.kind == :space && (!leading || entry.column > 1)
+            result << " "
+          else
+            result << segment_text(entry.start_offset, entry.end_offset, entry.text)
+          end
+        end
+      end
+
+      def segment_text(start_offset, end_offset, fallback)
+        return fallback unless source
+        return "" if end_offset <= start_offset
+
+        fallback_text = fallback.dup.force_encoding(source.encoding)
+        segment = source.byteslice(start_offset, end_offset - start_offset)
+        return fallback_text unless segment
+        segment = segment.dup.force_encoding(source.encoding)
+        return fallback_text if segment != fallback_text
+
+        segment
+      end
+    end
+  end
+end

data/lib/milk_tea/core/cst_builder.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module MilkTea
+  class CSTBuilder
+    def self.build(source, path: nil)
+      new(source, path:).build
+    end
+
+    def initialize(source, path: nil)
+      @source = source
+      @path = path
+    end
+
+    def build
+      lexed = Lexer.lex_with_trivia(@source, path: @path)
+      CST::SourceFile.new(source: @source, tokens: lexed.tokens, trivia: lexed.trivia)
+    end
+  end
+end