From eb5d9697994d75fb086e625e0331262f8b4610f9 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Sat, 14 Mar 2026 07:49:16 +0800 Subject: [PATCH] feat: implement iterator. --- apyds/chain_t.py | 17 ++++++ apyds/ds.cc | 39 +++++++++++++ apyds/search_t.py | 17 ++++++ atsds/ds.cc | 39 +++++++++++++ atsds/index.mts | 46 +++++++++++++++ include/ds/chain.hh | 5 ++ include/ds/generator.hh | 124 ++++++++++++++++++++++++++++++++++++++++ include/ds/search.hh | 5 ++ src/chain.cc | 66 ++++++++++++--------- src/search.cc | 52 ++++++++++------- tests/test_chain.cc | 14 +++++ tests/test_chain.mjs | 14 +++++ tests/test_chain.py | 13 +++++ tests/test_search.cc | 14 +++++ tests/test_search.mjs | 14 +++++ tests/test_search.py | 13 +++++ 16 files changed, 445 insertions(+), 47 deletions(-) create mode 100644 include/ds/generator.hh diff --git a/apyds/chain_t.py b/apyds/chain_t.py index 33d99447..f5b565e5 100644 --- a/apyds/chain_t.py +++ b/apyds/chain_t.py @@ -78,3 +78,20 @@ def execute(self, callback: typing.Callable[[Rule], bool]) -> int: The number of rules processed. """ return self._chain.execute(lambda candidate: callback(Rule(candidate.clone()))) + + def __iter__(self) -> typing.Iterator[Rule]: + """Iterate over inferred rules. + + Returns: + An iterator over Rule objects. + + Example: + >>> for rule in chain: + ... print(rule) + """ + iterator = self._chain.iter() + while True: + candidate = iterator.next() + if candidate is None: + break + yield Rule(candidate.clone()) diff --git a/apyds/ds.cc b/apyds/ds.cc index 42b86f99..142914fa 100644 --- a/apyds/ds.cc +++ b/apyds/ds.cc @@ -1,11 +1,37 @@ #include #include +#include #include #include #include namespace py = pybind11; +class Iterator { + public: + explicit Iterator(ds::generator _generator) : generator(std::move(_generator)), initialized(false), iterator(nullptr) { } + + ds::rule_t* next() { + if (initialized) { + ++*iterator; + } else { + iterator = std::make_unique(generator.begin()); + initialized = true; + } + if (*iterator == nullptr) { + return nullptr; + } + ds::rule_t* result = **iterator; + return result; + } + + private: + ds::generator generator; + bool initialized; + using iterator_t = decltype(generator.begin()); + std::unique_ptr iterator; +}; + template auto from_string(const std::string_view& string, int buffer_size) -> std::unique_ptr { auto result = reinterpret_cast(operator new(buffer_size)); @@ -169,6 +195,11 @@ PYBIND11_MODULE(_ds, m, py::mod_gil_not_used()) { search_t.def("reset", &ds::search_t::reset); search_t.def("add", &ds::search_t::add); search_t.def("execute", &ds::search_t::execute); + search_t.def( + "iter", + [](ds::search_t& self) { return Iterator(std::move(self.iterator())); }, + py::keep_alive<0, 1>() + ); auto chain_t = py::class_(m, "Chain"); chain_t.def(py::init()); @@ -177,4 +208,12 @@ PYBIND11_MODULE(_ds, m, py::mod_gil_not_used()) { chain_t.def("reset", &ds::chain_t::reset); chain_t.def("add", &ds::chain_t::add); chain_t.def("execute", &ds::chain_t::execute); + chain_t.def( + "iter", + [](ds::chain_t& self) { return Iterator(std::move(self.iterator())); }, + py::keep_alive<0, 1>() + ); + + auto iterator_t = py::class_(m, "Iterator"); + iterator_t.def("next", &Iterator::next, py::return_value_policy::reference_internal); } diff --git a/apyds/search_t.py b/apyds/search_t.py index edad6492..5b81bb28 100644 --- a/apyds/search_t.py +++ b/apyds/search_t.py @@ -77,3 +77,20 @@ def execute(self, callback: typing.Callable[[Rule], bool]) -> int: The number of rules processed. """ return self._search.execute(lambda candidate: callback(Rule(candidate.clone()))) + + def __iter__(self) -> typing.Iterator[Rule]: + """Iterate over inferred rules. + + Returns: + An iterator over Rule objects. + + Example: + >>> for rule in search: + ... print(rule) + """ + iterator = self._search.iter() + while True: + candidate = iterator.next() + if candidate is None: + break + yield Rule(candidate.clone()) diff --git a/atsds/ds.cc b/atsds/ds.cc index 86a91b8a..03bbe295 100644 --- a/atsds/ds.cc +++ b/atsds/ds.cc @@ -1,10 +1,36 @@ #include #include +#include #include #include namespace em = emscripten; +class Iterator { + public: + explicit Iterator(ds::generator _generator) : generator(std::move(_generator)), initialized(false), iterator(nullptr) { } + + ds::rule_t* next() { + if (initialized) { + ++*iterator; + } else { + iterator = std::make_unique(generator.begin()); + initialized = true; + } + if (*iterator == nullptr) { + return nullptr; + } + ds::rule_t* result = **iterator; + return result; + } + + private: + ds::generator generator; + bool initialized; + using iterator_t = decltype(generator.begin()); + std::unique_ptr iterator; +}; + // 由于embind的限制，这里无法使用string_view。 // 为了保持一致性，一律使用复制。 template @@ -134,6 +160,10 @@ auto search_execute(ds::search_t* search, const em::val& callback) -> ds::length return search->execute([&callback](ds::rule_t* candidate) -> bool { return callback(candidate, em::allow_raw_pointers()).as(); }); } +auto search_iter(ds::search_t* search) -> std::unique_ptr { + return std::make_unique(std::move(search->iterator())); +} + auto chain_add(ds::chain_t* chain, const std::string& text) -> bool { return chain->add(text); } @@ -142,6 +172,10 @@ auto chain_execute(ds::chain_t* chain, const em::val& callback) -> ds::length_t return chain->execute([&callback](ds::rule_t* candidate) -> bool { return callback(candidate, em::allow_raw_pointers()).as(); }); } +auto chain_iter(ds::chain_t* chain) -> std::unique_ptr { + return std::make_unique(std::move(chain->iterator())); +} + EMSCRIPTEN_BINDINGS(ds) { em::register_vector("Buffer"); @@ -195,6 +229,7 @@ EMSCRIPTEN_BINDINGS(ds) { // 因为embind的限制，这里无法使用string_view和function。 search_t.function("add", &search_add, em::allow_raw_pointers()); search_t.function("execute", &search_execute, em::allow_raw_pointers()); + search_t.function("iter", &search_iter, em::return_value_policy::take_ownership()); auto chain_t = em::class_("Chain"); chain_t.constructor(); @@ -204,4 +239,8 @@ EMSCRIPTEN_BINDINGS(ds) { // 因为 embind 的限制，这里无法使用 string_view 和 function。 chain_t.function("add", &chain_add, em::allow_raw_pointers()); chain_t.function("execute", &chain_execute, em::allow_raw_pointers()); + chain_t.function("iter", &chain_iter, em::return_value_policy::take_ownership()); + + auto iterator_t = em::class_("Iterator"); + iterator_t.function("next", &Iterator::next, em::return_value_policy::reference()); } diff --git a/atsds/index.mts b/atsds/index.mts index a40a36c3..2fc9f723 100644 --- a/atsds/index.mts +++ b/atsds/index.mts @@ -620,6 +620,29 @@ export class Search { return callback(new Rule(candidate).copy()); }); } + + /** + * Iterate over inferred rules. + * + * @returns An iterator over Rule objects. + * + * @example + * ```typescript + * for (const rule of search) { + * console.log(rule.toString()); + * } + * ``` + */ + *[Symbol.iterator](): Iterator { + const iterator = this._search.iter(); + while (true) { + const candidate = iterator.next(); + if (candidate === null) { + break; + } + yield new Rule(candidate); + } + } } /** @@ -697,4 +720,27 @@ export class Chain { return callback(new Rule(candidate).copy()); }); } + + /** + * Iterate over inferred rules. + * + * @returns An iterator over Rule objects. + * + * @example + * ```typescript + * for (const rule of chain) { + * console.log(rule.toString()); + * } + * ``` + */ + *[Symbol.iterator](): Iterator { + const iterator = this._chain.iter(); + while (true) { + const candidate = iterator.next(); + if (candidate === null) { + break; + } + yield new Rule(candidate); + } + } } diff --git a/include/ds/chain.hh b/include/ds/chain.hh index c2625fcb..eb0f540b 100644 --- a/include/ds/chain.hh +++ b/include/ds/chain.hh @@ -6,6 +6,7 @@ #include #include +#include #include namespace ds { @@ -70,6 +71,10 @@ namespace ds { /// @return 搜索到新的结果的数量。 /// @note 如果回调函数返回false，则继续搜索；如果回调函数返回true，则停止搜索。 length_t execute(const std::function& callback); + + /// @brief 执行一轮搜索操作，以生成器方式迭代所有匹配的规则。 + /// @return 生成器，每次迭代返回一个匹配的规则指针。 + generator iterator(); }; } // namespace ds diff --git a/include/ds/generator.hh b/include/ds/generator.hh new file mode 100644 index 00000000..8e5a3ba6 --- /dev/null +++ b/include/ds/generator.hh @@ -0,0 +1,124 @@ +#ifndef DS_GENERATOR_HH +#define DS_GENERATOR_HH + +#include +#include +#include + +namespace ds { + template + class _generator_promise; + + template + class _generator { + using promise_type = _generator_promise; + using handle_type = std::coroutine_handle; + handle_type handle_; + + public: + struct iterator { + handle_type h_; + + T& operator*() const noexcept { + return h_.promise().value_; + } + iterator& operator++() { + h_.resume(); + if (h_.done()) { + h_ = nullptr; + } + return *this; + } + bool operator==(std::nullptr_t) const noexcept { + return h_ == nullptr; + } + bool operator!=(std::nullptr_t) const noexcept { + return h_ != nullptr; + } + }; + + _generator(_generator&& other) noexcept : handle_(other.handle_) { + other.handle_ = nullptr; + } + _generator& operator=(_generator&& other) noexcept { + if (this != &other) { + if (handle_) { + handle_.destroy(); + } + handle_ = other.handle_; + other.handle_ = nullptr; + } + return *this; + } + + ~_generator() { + if (handle_) { + handle_.destroy(); + } + } + + iterator begin() { + if (!handle_) { + return {nullptr}; + } + handle_.resume(); + if (handle_.done()) { + return {nullptr}; + } + return {handle_}; + } + std::nullptr_t end() { + return nullptr; + } + + private: + _generator() noexcept : handle_(nullptr) { } + explicit _generator(handle_type h) noexcept : handle_(h) { } + friend class _generator_promise; + }; + + template + class _generator_promise { + public: + T value_; + + _generator get_return_object() noexcept { + return _generator(std::coroutine_handle<_generator_promise>::from_promise(*this)); + } + std::suspend_always initial_suspend() noexcept { + return {}; + } + std::suspend_always final_suspend() noexcept { + return {}; + } + std::suspend_always yield_value(T& value) noexcept { + value_ = value; + return {}; + } + std::suspend_always yield_value(T&& value) noexcept { + value_ = std::move(value); + return {}; + } + void return_void() noexcept { } + void unhandled_exception() { + std::terminate(); + } + }; + +#if defined(__cpp_lib_generator) && 0 + template + using generator = std::generator; +#else + template + using generator = _generator; +#endif +} // namespace ds + +namespace std { + template + struct coroutine_traits, Args...> { + using promise_type = ds::_generator_promise; + }; +} // namespace std + +#endif diff --git a/include/ds/search.hh b/include/ds/search.hh index 9ebd5f8a..90c1a09b 100644 --- a/include/ds/search.hh +++ b/include/ds/search.hh @@ -6,6 +6,7 @@ #include #include +#include #include namespace ds { @@ -65,6 +66,10 @@ namespace ds { /// @return 搜索到新的结果的数量。 /// @note 如果回调函数返回false，则继续搜索；如果回调函数返回true，则停止搜索。 length_t execute(const std::function& callback); + + /// @brief 执行一轮搜索操作，以生成器方式迭代所有匹配的规则。 + /// @return 生成器，每次迭代返回一个匹配的规则指针。 + generator iterator(); }; } // namespace ds diff --git a/src/chain.cc b/src/chain.cc index 4080e2b3..fdc9936c 100644 --- a/src/chain.cc +++ b/src/chain.cc @@ -63,28 +63,40 @@ namespace ds { } } - length_t chain_t::execute(const std::function& callback) { + ds::generator chain_t::iterator() { std::set, less_t> temp_facts; std::set, less_t> temp_rules; - bool break_all = false; + // RAII guard，确保无论是否提前退出，清理代码都会执行 + struct guard_t { + std::function cleanup; + ~guard_t() { + cleanup(); + } + } guard{[&]() { + ++current_cycle; + if (!temp_facts.empty()) { + last_fact_cycle = current_cycle; + } + for (auto it = temp_facts.begin(); it != temp_facts.end();) { + auto node = temp_facts.extract(it++); + facts.emplace(std::move(node.value()), current_cycle); + } + }}; - std::function chain_recursive; - chain_recursive = [&](rule_t* rule, rule_t* workspace, std::byte* tail) -> void { + auto chain_recursive = [&](auto& self, rule_t* rule, rule_t* workspace, std::byte* tail) -> ds::generator { if (rule->premises_count() == 0) { if (rule->data_size() > limit_size) { - return; + co_return; } auto new_fact = std::unique_ptr(reinterpret_cast(operator new(rule->data_size()))); memcpy(new_fact->head(), rule->head(), rule->data_size()); if (facts.find(new_fact) != facts.end() || temp_facts.find(new_fact) != temp_facts.end()) { - return; + co_return; } temp_facts.emplace(std::move(new_fact)); - if (callback(rule)) { - break_all = true; - } - return; + co_yield rule; + co_return; } else { do { if (rule->data_size() > limit_size) { @@ -96,9 +108,7 @@ namespace ds { break; } temp_rules.emplace(std::move(new_rule)); - if (callback(rule)) { - break_all = true; - } + co_yield rule; } while (false); } @@ -107,7 +117,9 @@ namespace ds { if (!workspace->valid()) { continue; } - chain_recursive(workspace, reinterpret_cast(workspace->tail()), tail); + for (auto yielded : self(self, workspace, reinterpret_cast(workspace->tail()), tail)) { + co_yield yielded; + } } }; @@ -116,21 +128,23 @@ namespace ds { continue; } - chain_recursive(rule.get(), buffer.get(), reinterpret_cast(buffer.get()) + buffer_size); - - if (break_all) { - break; + for (auto yielded : + chain_recursive(chain_recursive, rule.get(), buffer.get(), reinterpret_cast(buffer.get()) + buffer_size)) { + co_yield yielded; } } - if (!break_all) { - done_cycle = current_cycle; - } - ++current_cycle; - length_t count = temp_rules.size() + temp_facts.size(); - for (auto it = temp_facts.begin(); it != temp_facts.end();) { - auto node = temp_facts.extract(it++); - facts.emplace(std::move(node.value()), current_cycle); + done_cycle = current_cycle; + co_return; + } + + length_t chain_t::execute(const std::function& callback) { + length_t count = 0; + for (auto* rule : iterator()) { + ++count; + if (callback(rule)) { + break; + } } return count; } diff --git a/src/search.cc b/src/search.cc index 47feaf62..4465c443 100644 --- a/src/search.cc +++ b/src/search.cc @@ -61,11 +61,28 @@ namespace ds { } } - length_t search_t::execute(const std::function& callback) { + ds::generator search_t::iterator() { std::set, less_t> temp_rules; std::set, less_t> temp_facts; - bool break_all = false; + // RAII guard，确保无论是否提前退出，清理代码都会执行 + struct guard_t { + std::function cleanup; + ~guard_t() { + cleanup(); + } + } guard{[&]() { + ++current_cycle; + for (auto it = temp_rules.begin(); it != temp_rules.end();) { + auto node = temp_rules.extract(it++); + rules.emplace(std::move(node.value()), current_cycle); + } + for (auto it = temp_facts.begin(); it != temp_facts.end();) { + auto node = temp_facts.extract(it++); + facts.emplace(std::move(node.value()), current_cycle); + } + }}; + for (auto& [rule, rules_cycle] : rules) { for (auto& [fact, facts_cycle] : facts) { if (rules_cycle <= done_cycle && facts_cycle <= done_cycle) { @@ -95,28 +112,21 @@ namespace ds { memcpy(new_fact.get(), buffer.get(), buffer->data_size()); temp_facts.emplace(std::move(new_fact)); } - if (callback(buffer.get())) { - break_all = true; - break; - } - } - if (break_all) { - break; + co_yield buffer.get(); } } - if (!break_all) { - done_cycle = current_cycle; - } - ++current_cycle; - length_t count = temp_rules.size() + temp_facts.size(); - for (auto it = temp_rules.begin(); it != temp_rules.end();) { - auto node = temp_rules.extract(it++); - rules.emplace(std::move(node.value()), current_cycle); - } - for (auto it = temp_facts.begin(); it != temp_facts.end();) { - auto node = temp_facts.extract(it++); - facts.emplace(std::move(node.value()), current_cycle); + done_cycle = current_cycle; + co_return; + } + + length_t search_t::execute(const std::function& callback) { + length_t count = 0; + for (auto* rule : iterator()) { + ++count; + if (callback(rule)) { + break; + } } return count; } diff --git a/tests/test_chain.cc b/tests/test_chain.cc index 3c9af4bb..8c36a7cd 100644 --- a/tests/test_chain.cc +++ b/tests/test_chain.cc @@ -146,3 +146,17 @@ TEST_F(TestChain, execute_exceed_by_too_many_premises) { EXPECT_TRUE(chain->add("eeeee")); EXPECT_EQ(chain->execute([](ds::rule_t* rule) { return false; }), 1); } + +TEST_F(TestChain, iterator) { + chain->add("a"); + chain->add("b"); + chain->add("a b c"); + const char* expected[] = {"b\n----\nc\n", "----\nc\n"}; + int count = 0; + for (auto rule : chain->iterator()) { + EXPECT_LT(count, 2); + EXPECT_STREQ(ds::rule_to_text(rule, limit_size).get(), expected[count]); + ++count; + } + EXPECT_EQ(count, 2); +} diff --git a/tests/test_chain.mjs b/tests/test_chain.mjs index 87629eb3..f7d05261 100644 --- a/tests/test_chain.mjs +++ b/tests/test_chain.mjs @@ -134,3 +134,17 @@ test("execute_exceed_by_too_many_premises", () => { expect(newChain.add("eeeee")).toBe(true); expect(newChain.execute((rule) => false)).toBe(1); }); + +test("iterator", () => { + chain.add("a"); + chain.add("b"); + chain.add("a b c"); + const expected = ["b\n----\nc\n", "----\nc\n"]; + let count = 0; + for (const rule of chain) { + expect(count).toBeLessThan(expected.length); + expect(rule.toString()).toBe(expected[count]); + count++; + } + expect(count).toBe(expected.length); +}); diff --git a/tests/test_chain.py b/tests/test_chain.py index b2986d57..98e6da45 100644 --- a/tests/test_chain.py +++ b/tests/test_chain.py @@ -140,3 +140,16 @@ def test_execute_exceed_by_too_many_premises() -> None: assert chain.add("ddddd") assert chain.add("eeeee") assert chain.execute(lambda rule: False) == 1 + + +def test_iterator(chain: apyds.Chain) -> None: + chain.add("a") + chain.add("b") + chain.add("a b c") + expected = ["b\n----\nc\n", "----\nc\n"] + count = 0 + for rule in chain: + assert count < len(expected) + assert str(rule) == expected[count] + count += 1 + assert count == len(expected) diff --git a/tests/test_search.cc b/tests/test_search.cc index a0fb95e4..6bc416de 100644 --- a/tests/test_search.cc +++ b/tests/test_search.cc @@ -102,3 +102,17 @@ TEST_F(TestSearch, execute_exceed) { auto count = search->execute([](ds::rule_t* rule) { return false; }); EXPECT_EQ(count, 0); } + +TEST_F(TestSearch, iterator) { + search->add("a"); + search->add("b"); + search->add("a b c"); + const char* expected[] = {"b\n----\nc\n"}; + int count = 0; + for (auto rule : search->iterator()) { + EXPECT_LT(count, 1); + EXPECT_STREQ(ds::rule_to_text(rule, limit_size).get(), expected[count]); + ++count; + } + EXPECT_EQ(count, 1); +} diff --git a/tests/test_search.mjs b/tests/test_search.mjs index eb03c574..411ba43d 100644 --- a/tests/test_search.mjs +++ b/tests/test_search.mjs @@ -88,3 +88,17 @@ test("execute_exceed", () => { const count = search.execute((rule) => false); expect(count).toBe(0); }); + +test("iterator", () => { + search.add("a"); + search.add("b"); + search.add("a b c"); + const expected = ["b\n----\nc\n"]; + let count = 0; + for (const rule of search) { + expect(count).toBeLessThan(expected.length); + expect(rule.toString()).toBe(expected[count]); + count++; + } + expect(count).toBe(expected.length); +}); diff --git a/tests/test_search.py b/tests/test_search.py index a2d76ccb..bbfcb67e 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -94,3 +94,16 @@ def test_execute_exceed(search: apyds.Search) -> None: assert search.add("(2 a-very-long-fact-that-exceeds-half-of-the-limit-size)") count = search.execute(lambda rule: False) assert count == 0 + + +def test_iterator(search: apyds.Search) -> None: + search.add("a") + search.add("b") + search.add("a b c") + expected = ["b\n----\nc\n"] + count = 0 + for rule in search: + assert count < len(expected) + assert str(rule) == expected[count] + count += 1 + assert count == len(expected)