diff --git a/src/ast/block_statement.h b/src/ast/block_statement.h index 4ee0806582..3c2a26874e 100644 --- a/src/ast/block_statement.h +++ b/src/ast/block_statement.h @@ -42,6 +42,13 @@ class BlockStatement : public Statement { statements_.push_back(std::move(stmt)); } + /// Insert a statement to the block + /// @param index the index to insert at + /// @param stmt the statement to insert + void insert(size_t index, std::unique_ptr stmt) { + statements_.insert(statements_.begin() + index, std::move(stmt)); + } + /// @returns true if the block is empty bool empty() const { return statements_.empty(); } /// @returns the number of statements directly in the block diff --git a/src/ast/block_statement_test.cc b/src/ast/block_statement_test.cc index f0a6cbd263..2d040d3051 100644 --- a/src/ast/block_statement_test.cc +++ b/src/ast/block_statement_test.cc @@ -38,6 +38,27 @@ TEST_F(BlockStatementTest, Creation) { EXPECT_EQ(b[0], ptr); } +TEST_F(BlockStatementTest, Creation_WithInsert) { + auto s1 = std::make_unique(); + auto s2 = std::make_unique(); + auto s3 = std::make_unique(); + auto* p1 = s1.get(); + auto* p2 = s2.get(); + auto* p3 = s3.get(); + + BlockStatement b; + b.insert(0, std::move(s1)); + b.insert(0, std::move(s2)); + b.insert(1, std::move(s3)); + + // |b| should contain s2, s3, s1 + + ASSERT_EQ(b.size(), 3u); + EXPECT_EQ(b[0], p2); + EXPECT_EQ(b[1], p3); + EXPECT_EQ(b[2], p1); +} + TEST_F(BlockStatementTest, Creation_WithSource) { BlockStatement b(Source{20, 2}); auto src = b.source(); diff --git a/src/ast/module.h b/src/ast/module.h index d17a0cb44c..7734026621 100644 --- a/src/ast/module.h +++ b/src/ast/module.h @@ -57,6 +57,9 @@ class Module { /// @returns the global variables for the module const VariableList& global_variables() const { return global_variables_; } + /// @returns the global variables for the module + VariableList& global_variables() { return global_variables_; } + /// Adds an entry point to the module /// @param ep the entry point to add void AddEntryPoint(std::unique_ptr ep) {