Skip to content

Commit 079e54e

Browse files
authored
Add Module API to query for the number of methods (#9441)
1 parent 256dfa0 commit 079e54e

File tree

3 files changed

+21
-0
lines changed

3 files changed

+21
-0
lines changed

extension/module/module.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,11 @@ runtime::Error Module::load(const runtime::Program::Verification verification) {
164164
return runtime::Error::Ok;
165165
}
166166

167+
runtime::Result<size_t> Module::num_methods() {
168+
ET_CHECK_OK_OR_RETURN_ERROR(load());
169+
return program_->num_methods();
170+
}
171+
167172
runtime::Result<std::unordered_set<std::string>> Module::method_names() {
168173
ET_CHECK_OK_OR_RETURN_ERROR(load());
169174
const auto method_count = program_->num_methods();

extension/module/module.h

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,14 @@ class Module {
138138
return program_;
139139
}
140140

141+
/**
142+
* Get the number of methods available in the loaded program.
143+
*
144+
* @returns A Result object containing either the number of methods available
145+
* or an error to indicate failure.
146+
*/
147+
runtime::Result<size_t> num_methods();
148+
141149
/**
142150
* Get a list of method names available in the loaded program.
143151
* Loads the program and method if needed.

extension/module/test/module_test.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,14 @@ TEST_F(ModuleTest, TestMethodNames) {
6969
EXPECT_EQ(method_names.get(), std::unordered_set<std::string>{"forward"});
7070
}
7171

72+
TEST_F(ModuleTest, TestNumMethods) {
73+
Module module(model_path_);
74+
75+
const auto num_methods = module.num_methods();
76+
EXPECT_EQ(num_methods.error(), Error::Ok);
77+
EXPECT_EQ(num_methods.get(), 1);
78+
}
79+
7280
TEST_F(ModuleTest, TestNonExistentMethodNames) {
7381
Module module("/path/to/nonexistent/file.pte");
7482

0 commit comments

Comments
 (0)