Generalize InitialBlockVisitor

- use a template
- add same visitor for BREAKPOINT block
- add tests

Author: Goran Jelic-Cizmek

src/nmodl/main.cpp

@@ -28,7 +28,6 @@
 #include "visitors/cvode_visitor.hpp"
 #include "visitors/function_callpath_visitor.hpp"
 #include "visitors/global_var_visitor.hpp"
-#include "visitors/initial_block_visitor.hpp"
 #include "visitors/implicit_argument_visitor.hpp"
 #include "visitors/indexedname_visitor.hpp"
 #include "visitors/inline_visitor.hpp"
@@ -39,6 +38,7 @@
 #include "visitors/localize_visitor.hpp"
 #include "visitors/longitudinal_diffusion_visitor.hpp"
 #include "visitors/loop_unroll_visitor.hpp"
+#include "visitors/merge_top_level_blocks_visitor.hpp"
 #include "visitors/neuron_solve_visitor.hpp"
 #include "visitors/nmodl_visitor.hpp"
 #include "visitors/perf_visitor.hpp"
@@ -347,10 +347,19 @@ int run_nmodl(int argc, const char* argv[]) {
         /// merge all INITIAL blocks into one (this needs to run before SymtabVisitor)
         {
             logger->info("Running INITIAL block merge visitor");
-            MergeInitialBlocksVisitor().visit_program(*ast);
+            MergeTopLevelBlocksVisitor<ast::InitialBlock, ast::AstNodeType::INITIAL_BLOCK>()
+                .visit_program(*ast);
             ast_to_nmodl(*ast, filepath("merge_initial_block"));
         }
 
+        /// merge all BREAKPOINT blocks into one (this needs to run before SymtabVisitor)
+        {
+            logger->info("Running BREAKPOINT block merge visitor");
+            MergeTopLevelBlocksVisitor<ast::BreakpointBlock, ast::AstNodeType::BREAKPOINT_BLOCK>()
+                .visit_program(*ast);
+            ast_to_nmodl(*ast, filepath("merge_breakpoint_block"));
+        }
+
         /// construct symbol table
         {
             logger->info("Running symtab visitor");

src/nmodl/visitors/CMakeLists.txt

@@ -28,7 +28,6 @@ add_library(
   neuron_solve_visitor.cpp
   perf_visitor.cpp
   rename_visitor.cpp
-  initial_block_visitor.cpp
   rename_function_arguments.cpp
   semantic_analysis_visitor.cpp
   solve_block_visitor.cpp

src/nmodl/visitors/initial_block_visitor.cpp

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2025 EPFL.
+ * See the top-level LICENSE file for details.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+/**
+ * \file
+ * \brief \copybrief nmodl::visitor::MergeTopLevelBlocksVisitor
+ */
+
+#include "visitors/ast_visitor.hpp"
+
+#include "visitors/visitor_utils.hpp"
+
+namespace nmodl {
+namespace visitor {
+
+/**
+ * \addtogroup visitor_classes
+ * \{
+ */
+
+/**
+ * \class MergeTopLevelBlocksVisitor
+ * \brief Visitor which merges given top-level blocks into one
+ */
+template <typename ast_class, ast::AstNodeType ast_type>
+class MergeTopLevelBlocksVisitor: public AstVisitor {
+  public:
+    MergeTopLevelBlocksVisitor() = default;
+    void visit_program(ast::Program& node) override {
+        // check if there is > 1 blocks in total
+        if (collect_nodes(node, {ast_type}).size() <= 1) {
+            return;
+        }
+
+        // collect all top-level blocks in the program
+        const auto& blocks = node.get_blocks();
+        logger->info("Collecting blocks of type {}", typeid(ast_class).name());
+
+        // collect all statements from top-level blocks, and the blocks themselves
+        ast::StatementVector statements;
+        std::unordered_set<ast::Node*> blocks_to_delete;
+        for (auto& block: blocks) {
+            auto temp_block = std::dynamic_pointer_cast<ast_class>(block);
+            // check if it's the correct type
+            if (temp_block) {
+                auto statement_block = temp_block->get_statement_block();
+                // if block is not empty, copy statements into vector
+                if (statement_block) {
+                    for (const auto& statement: statement_block->get_statements()) {
+                        statements.push_back(statement);
+                    }
+                }
+                blocks_to_delete.insert(block.get());
+            }
+        }
+
+        // insert new top-level block which has the above statements
+        auto new_block = ast::StatementBlock(statements);
+        auto new_initial_block = ast_class(new_block.clone());
+        node.emplace_back_node(new_initial_block.clone());
+
+        // delete all of the previously-found top-level blocks
+        node.erase_node(blocks_to_delete);
+    }
+};
+
+/** \} */  // end of visitor_classes
+
+}  // namespace visitor
+}  // namespace nmodl

src/nmodl/visitors/initial_block_visitor.hpp

@@ -48,7 +48,6 @@ add_executable(
   visitor/global_to_range.cpp
   visitor/implicit_argument.cpp
   visitor/inline.cpp
-  visitor/initial_block.cpp
   visitor/json.cpp
   visitor/kinetic_block.cpp
   visitor/localize.cpp
@@ -57,6 +56,7 @@ add_executable(
   visitor/local_to_assigned.cpp
   visitor/lookup.cpp
   visitor/loop_unroll.cpp
+  visitor/merge_top_level_blocks.cpp
   visitor/misc.cpp
   visitor/neuron_solve.cpp
   visitor/nmodl.cpp