[ruby/prism] Provide API for visiting in C

https://github.com/ruby/prism/commit/537947aa5c

2 files changed, 117 insertions, 0 deletions

diff --git a/prism/node.h b/prism/node.h
index 76eb720978..a001b4a9e4 100644
--- a/prism/node.h
+++ b/prism/node.h
@@ -82,4 +82,66 @@ PRISM_EXPORTED_FUNCTION void pm_node_memsize(pm_node_t *node, pm_memsize_t *mems
  */
 PRISM_EXPORTED_FUNCTION const char * pm_node_type_to_str(pm_node_type_t node_type);
 
+/**
+ * Visit each of the nodes in this subtree using the given visitor callback. The
+ * callback function will be called for each node in the subtree. If it returns
+ * false, then that node's children will not be visited. If it returns true,
+ * then the children will be visited. The data parameter is treated as an opaque
+ * pointer and is passed to the visitor callback for consumers to use as they
+ * see fit.
+ *
+ * As an example:
+ *
+ * ```c
+ * #include "prism.h"
+ *
+ * bool visit(const pm_node_t *node, void *data) {
+ *     size_t *indent = (size_t *) data;
+ *     for (size_t i = 0; i < *indent * 2; i++) putc(' ', stdout);
+ *     printf("%s\n", pm_node_type_to_str(node->type));
+ *
+ *     size_t next_indent = *indent + 1;
+ *     size_t *next_data = &next_indent;
+ *     pm_visit_child_nodes(node, visit, next_data);
+ *
+ *     return false;
+ * }
+ *
+ * int main(void) {
+ *     const char *source = "1 + 2; 3 + 4";
+ *     size_t size = strlen(source);
+ *
+ *     pm_parser_t parser;
+ *     pm_options_t options = { 0 };
+ *     pm_parser_init(&parser, (const uint8_t *) source, size, &options);
+ *
+ *     size_t indent = 0;
+ *     pm_node_t *node = pm_parse(&parser);
+ *
+ *     size_t *data = &indent;
+ *     pm_visit_node(node, visit, data);
+ *
+ *     pm_node_destroy(&parser, node);
+ *     pm_parser_free(&parser);
+ *     return EXIT_SUCCESS;
+ * }
+ * ```
+ *
+ * @param node The root node to start visiting from.
+ * @param visitor The callback to call for each node in the subtree.
+ * @param data An opaque pointer that is passed to the visitor callback.
+ */
+PRISM_EXPORTED_FUNCTION void pm_visit_node(const pm_node_t *node, bool (*visitor)(const pm_node_t *node, void *data), void *data);
+
+/**
+ * Visit the children of the given node with the given callback. This is the
+ * default behavior for walking the tree that is called from pm_visit_node if
+ * the callback returns true.
+ *
+ * @param node The node to visit the children of.
+ * @param visitor The callback to call for each child node.
+ * @param data An opaque pointer that is passed to the visitor callback.
+ */
+PRISM_EXPORTED_FUNCTION void pm_visit_child_nodes(const pm_node_t *node, bool (*visitor)(const pm_node_t *node, void *data), void *data);
+
 #endif
diff --git a/prism/templates/src/node.c.erb b/prism/templates/src/node.c.erb
index 419cb48bf0..d689c2c66b 100644
--- a/prism/templates/src/node.c.erb
+++ b/prism/templates/src/node.c.erb
@@ -192,6 +192,61 @@ pm_node_type_to_str(pm_node_type_t node_type)
     return "";
 }
 
+/**
+ * Visit each of the nodes in this subtree using the given visitor callback. The
+ * callback function will be called for each node in the subtree. If it returns
+ * false, then that node's children will not be visited. If it returns true,
+ * then the children will be visited. The data parameter is treated as an opaque
+ * pointer and is passed to the visitor callback for consumers to use as they
+ * see fit.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_visit_node(const pm_node_t *node, bool (*visitor)(const pm_node_t *node, void *data), void *data) {
+    if (visitor(node, data)) pm_visit_child_nodes(node, visitor, data);
+}
+
+/**
+ * Visit the children of the given node with the given callback. This is the
+ * default behavior for walking the tree that is called from pm_visit_node if
+ * the callback returns true.
+ */
+PRISM_EXPORTED_FUNCTION void
+pm_visit_child_nodes(const pm_node_t *node, bool (*visitor)(const pm_node_t *node, void *data), void *data) {
+    switch (PM_NODE_TYPE(node)) {
+        <%- nodes.each do |node| -%>
+        <%- if (fields = node.fields.select { |field| field.is_a?(Prism::NodeField) || field.is_a?(Prism::OptionalNodeField) || field.is_a?(Prism::NodeListField) }).any? -%>
+        case <%= node.type %>: {
+            const pm_<%= node.human %>_t *cast = (const pm_<%= node.human %>_t *) node;
+            <%- fields.each do |field| -%>
+
+            // Visit the <%= field.name %> field
+            <%- case field -%>
+            <%- when Prism::NodeField -%>
+            pm_visit_node((const pm_node_t *) cast-><%= field.name %>, visitor, data);
+            <%- when Prism::OptionalNodeField -%>
+            if (cast-><%= field.name %> != NULL) {
+                pm_visit_node((const pm_node_t *) cast-><%= field.name %>, visitor, data);
+            }
+            <%- when Prism::NodeListField -%>
+            const pm_node_list_t *<%= field.name %> = &cast-><%= field.name %>;
+            for (size_t index = 0; index < <%= field.name %>->size; index++) {
+                pm_visit_node(<%= field.name %>->nodes[index], visitor, data);
+            }
+            <%- end -%>
+            <%- end -%>
+
+            break;
+        }
+        <%- else -%>
+        case <%= node.type %>:
+            break;
+        <%- end -%>
+        <%- end -%>
+        case PM_SCOPE_NODE:
+            break;
+    }
+}
+
 static void
 pm_dump_json_constant(pm_buffer_t *buffer, const pm_parser_t *parser, pm_constant_id_t constant_id) {
     const pm_constant_t *constant = pm_constant_pool_id_to_constant(&parser->constant_pool, constant_id);