[ruby/prism] Rewrite logic for CHECK_FIELD_KIND to improve readability

https://github.com/ruby/prism/commit/f731edcc26

2 files changed, 27 insertions, 23 deletions

diff --git a/prism/templates/lib/prism/node.rb.erb b/prism/templates/lib/prism/node.rb.erb
index 0afaa28310..4762963bf6 100644
--- a/prism/templates/lib/prism/node.rb.erb
+++ b/prism/templates/lib/prism/node.rb.erb
@@ -98,9 +98,6 @@ module Prism
       raise NoMethodError, "undefined method `type' for #{inspect}"
     end
   end
-  <%- if ENV["CHECK_FIELD_KIND"] -%>
-  CHECK_FIELD_KIND = ENV["CHECK_FIELD_KIND"]
-  <%- end -%>
   <%- nodes.each do |node| -%>
 
   <%- node.each_comment_line do |line| -%>
@@ -113,26 +110,8 @@ module Prism
       @newline = false
       @location = location
       <%- node.fields.each do |field| -%>
-      <%- if ENV["CHECK_FIELD_KIND"] -%>
-      <%- if field.respond_to?(:union_kind) && field.union_kind -%>
-      <%- case field -%>
-      <%- when Prism::NodeField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && ![<%= field.union_kind.join(', ') %>].include?(<%= field.name %>.class)
-      <%- when Prism::OptionalNodeField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && ![<%= field.union_kind.join(', ') %>, NilClass].include?(<%= field.name %>.class)
-      <%- when Prism::NodeListField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && !<%= field.name %>.all? { |n| [<%= field.union_kind.join(', ') %>].include?(n.class) }
-      <%- end -%>
-      <%- elsif field.respond_to?(:specific_kind) && field.specific_kind -%>
-      <%- case field -%>
-      <%- when Prism::NodeField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && !<%= field.name %>.is_a?(<%= field.specific_kind %>)
-      <%- when Prism::OptionalNodeField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && !<%= field.name %>.nil? && !<%= field.name %>.is_a?(<%= field.specific_kind %>)
-      <%- when Prism::NodeListField -%>
-      raise <%= field.name %>.inspect if CHECK_FIELD_KIND && !<%= field.name %>.all? { |n| n.is_a?(<%= field.specific_kind %>) }
-      <%- end -%>
-      <%- end -%>
+      <%- if Prism::CHECK_FIELD_KIND && field.respond_to?(:check_field_kind) -%>
+      raise <%= field.name %>.inspect unless <%= field.check_field_kind %>
       <%- end -%>
       @<%= field.name %> = <%= field.name %>
       <%- end -%>
diff --git a/prism/templates/template.rb b/prism/templates/template.rb
index 466718e17b..fd55d5228b 100755
--- a/prism/templates/template.rb
+++ b/prism/templates/template.rb
@@ -6,6 +6,7 @@ require "yaml"
 
 module Prism
   SERIALIZE_ONLY_SEMANTICS_FIELDS = ENV.fetch("PRISM_SERIALIZE_ONLY_SEMANTICS_FIELDS", false)
+  CHECK_FIELD_KIND = ENV.fetch("CHECK_FIELD_KIND", false)
 
   JAVA_BACKEND = ENV["PRISM_JAVA_BACKEND"] || "truffleruby"
   JAVA_STRING_TYPE = JAVA_BACKEND == "jruby" ? "org.jruby.RubySymbol" : "String"
@@ -123,6 +124,14 @@ module Prism
     def rbi_class
       "Prism::#{ruby_type}"
     end
+
+    def check_field_kind
+      if union_kind
+        "[#{union_kind.join(', ')}].include?(#{name}.class)"
+      else
+        "#{name}.is_a?(#{ruby_type})"
+      end
+    end
   end
 
   # This represents a field on a node that is itself a node and can be
@@ -141,6 +150,14 @@ module Prism
     def rbi_class
       "T.nilable(Prism::#{ruby_type})"
     end
+
+    def check_field_kind
+      if union_kind
+        "[#{union_kind.join(', ')}, NilClass].include?(#{name}.class)"
+      else
+        "#{name}.nil? || #{name}.is_a?(#{ruby_type})"
+      end
+    end
   end
 
   # This represents a field on a node that is a list of nodes. We pass them as
@@ -163,6 +180,14 @@ module Prism
     def java_type
       "#{super}[]"
     end
+
+    def check_field_kind
+      if union_kind
+        "#{name}.all? { |n| [#{union_kind.join(', ')}].include?(n.class) }"
+      else
+        "#{name}.all? { |n| n.is_a?(#{ruby_type}) }"
+      end
+    end
   end
 
   # This represents a field on a node that is the ID of a string interned