rework command_map a little

4 files changed, 248 insertions, 129 deletions

diff --git a/nihil/CMakeLists.txt b/nihil/CMakeLists.txt
index c84cd32..8aa4e27 100644
--- a/nihil/CMakeLists.txt
+++ b/nihil/CMakeLists.txt
@@ -29,6 +29,7 @@ target_sources(nihil
 	write_file.ccm
 
 	PRIVATE
+	command_map.cc
 	ensure_dir.cc
 	exec.cc
 	fd.cc
diff --git a/nihil/command_map.cc b/nihil/command_map.cc
new file mode 100644
index 0000000..1cc5081
--- /dev/null
+++ b/nihil/command_map.cc
@@ -0,0 +1,217 @@
+/*
+ * This source code is released into the public domain.
+ */
+
+module;
+
+#include <cstdio>
+#include <functional>
+#include <iostream>
+#include <map>
+#include <ranges>
+#include <string>
+#include <utility>
+
+module nihil;
+
+/*
+ * command_map represents a hierarchical list of commands.  At each level,
+ * a command is mapped to a handler, which can either be a function, in
+ * which case we execute the function, or another command_map, in which
+ * case we invoke the new map
+ */
+
+namespace nihil {
+
+/*
+ * The string tree we store our commands in.  This is sort of like a very
+ * basic hierarchical std::map.  Keys are provided as a range of values,
+ * typically from argv.
+ */
+
+template<typename T>
+struct string_tree_node final {
+	string_tree_node()
+		: _this_word("")
+	{
+	}
+
+	string_tree_node(std::string_view this_word)
+		: _this_word(this_word)
+	{
+	}
+
+	string_tree_node(std::string_view this_word, T value)
+		: _this_word(this_word)
+		, _value(std::move(value))
+	{
+	}
+
+	/*
+	 * Return a child node, or NULL if the child doesn't exist.
+	 */
+	auto get_child(this string_tree_node const &self,
+		       std::string_view child)
+		-> string_tree_node const *
+	{
+		if (auto it = self.children.find(std::string(child));
+		    it != self.children.end())
+			return &it->second;
+
+		return nullptr;
+	}
+
+	auto get_child(this string_tree_node &self,
+		       std::string_view child)
+		-> string_tree_node *
+	{
+		if (auto it = self.children.find(std::string(child));
+		    it != self.children.end())
+			return &it->second;
+
+		return nullptr;
+	}
+
+	/*
+	 * Return a child node if it exists, or insert a new empty node.
+	 */
+	auto get_or_create_child(this string_tree_node &self,
+				 std::string_view child)
+		-> string_tree_node *
+	{
+		if (auto ptr = self.get_child(child); ptr != nullptr)
+			return ptr;
+
+		auto [it, ok] = self.children.emplace(
+					child, string_tree_node(child));
+		return &it->second;
+	}
+
+	/*
+	 * Return this node's value.
+	 */
+	auto value(this string_tree_node const &self)
+		-> std::optional<T> const &
+	{
+		return self._value;
+	}
+
+	/*
+	 * Set this node's value.
+	 */
+	auto value(this string_tree_node &self, T value) -> void
+	{
+		self._value = std::move(value);
+	}
+
+private:
+	std::string _this_word;
+	std::optional<T> _value;
+	std::unordered_map<std::string, string_tree_node> children;
+};
+
+template<typename T>
+struct string_tree {
+	/*
+	 * Add a node to the tree.  Returns false if the node already exists.
+	 */
+	auto insert(this string_tree &self,
+		    std::ranges::range auto &&path,
+		    T value)
+		-> bool
+	{
+		auto *this_node = &self._root_node;
+
+		// Find the node for this key.
+		for (auto &&next : path) {
+			auto this_word = std::string_view(next);
+			this_node = this_node->get_or_create_child(this_word);
+		}
+
+		if (this_node->value()) {
+			// The value already exists.
+			return false;
+		}
+
+		// Set the new value.
+		this_node->value(std::move(value));
+		return true;
+	}
+
+	/*
+	 * Find a node in the tree.  Unlike insert(), iteration stops when
+	 * we find any node with a value.
+	 */
+	auto find(this string_tree const &self,
+		  std::ranges::range auto &&path)
+		-> std::optional<T>
+	{
+		auto *this_node = &self._root_node;
+{}
+		// Find the node for this key.
+		for (auto &&next : path) {
+			this_node = this_node->get_child(next);
+
+			if (this_node == nullptr)
+				// The node doesn't exist.
+				return {};
+
+			if (this_node->value())
+				// This node has a value, so return it.
+				return {this_node->value()};
+		}
+
+		// We didn't find a value, so the key was incomplete.
+		return {};
+	}
+
+private:
+	string_tree_node<T> _root_node;
+};
+
+/*
+ * The global command map.
+ */
+auto get_commands() -> string_tree<command_base *> & {
+	static auto commands = string_tree<command_base *>();
+	return commands;
+}
+
+auto register_command(std::string_view path, command_base *cmd) noexcept
+	-> void
+try {
+{
+	auto &commands = get_commands();
+	if (commands.insert(path | std::views::split(' '), cmd) == false) {
+		std::printf("command registration failed\n");
+		std::abort();
+	}
+}
+} catch (...) {
+	std::printf("command registration failed\n");
+	std::abort();
+}
+
+
+auto dispatch_command(int argc, char **argv) -> int
+{
+	if (argc == 0) {
+		std::print("not enough arguments\n");
+		return 1;
+	}
+
+	auto &commands = get_commands();
+	auto node = commands.find(std::span(argv, argv + argc));
+	if (node)
+		return (**node).invoke(argc, argv);
+
+	std::print("unknown command\n");
+	return 1;
+}
+
+void print_usage(std::string_view)
+{
+//	get_root_node().print_usage(std::string(prefix));
+}
+
+} // namespace nihil
diff --git a/nihil/command_map.ccm b/nihil/command_map.ccm
index 279c43e..979657b 100644
--- a/nihil/command_map.ccm
+++ b/nihil/command_map.ccm
@@ -24,130 +24,52 @@ import :usage_error;
 
 namespace nihil {
 
-template<typename Context>
-struct node;
-
-template<typename Context>
-auto get_root_node() noexcept -> node<Context> &;
-
-// Declare a global command and add it to the root node.
-export template<typename Context>
-struct command final {
-	template<typename F>
-	command(std::string_view path, F fn)
-	try : _func(std::move(fn))
+struct command_base {
+	command_base(std::string_view path)
+		: _path(path)
 	{
-		auto &node = get_root_node<Context>().create_node(path);
-		node.handler = this;
-	} catch (std::exception const &exc) {
-		std::cerr << "ERROR: failed to initialise command "
-			<< path << ": " << exc.what() << "\n";
-		std::abort();
 	}
 
-	auto invoke(Context const &ctx, int argc, char **argv) -> int
+	auto path(this command_base const &self) -> std::string_view
 	{
-		return std::invoke(_func, ctx, argc, argv);
+		return self._path;
 	}
 
+	virtual auto invoke(int argc, char **argv) -> int = 0;
+
 private:
-	std::function<int (Context const &, int, char **)> _func;
+	std::string_view _path;
 };
 
-// A node in the command hierarchy.
-template<typename Context>
-struct node {
-	// The command name of this node.
-	std::string name;
-
-	// Handler for this node.  May be null, which means this node has
-	// sub-commands but isn't a command itself.
-	command<Context> *handler = nullptr;
-
-	node(std::string name_)
-		: name(std::move(name_))
-	{}
+/*
+ * Register a command; used by command<>::command().
+ */
+auto register_command(std::string_view path, command_base *) noexcept -> void;
 
-	// Run the handler for this node.
-	auto invoke(Context const &ctx, int argc, char **argv) const -> int
+/*
+ * A command that can be invoked.  Instantiating a command adds this command
+ * to the global command table.  If an error occurs, the program will abort.
+ */
+export template<typename Fn>
+struct command final : command_base {
+	command(std::string_view path, Fn func) noexcept
+		: command_base(path)
+		, _func(func)
 	{
-		if (argc == 0)
-			throw usage_error("incomplete command");
-
-		// Look for a subcommand with argv[0].
-		auto it = commands.find(argv[0]);
-		if (it == commands.end())
-			throw usage_error("unknown command");
-
-		auto const &child = it->second;
-
-		// If the child has a handler, invoke it.
-		if (child.handler != nullptr)
-			return child.handler->invoke(ctx, argc, argv);
-
-		--argc;
-		++argv;
-
-		// Otherwise, continue recursing.
-		return child.invoke(ctx, argc, argv);
+		register_command(path, this);
 	}
 
-
-	// Create a new node under this one, or return it if it already exists.
-	// If path is empty, return this node.
-	auto create_node(std::string_view path) -> node&
+	auto invoke(int argc, char **argv) -> int override
 	{
-		auto child = next_word(&path);
-		if (child.empty())
-			return *this;
-
-		// If the child node doesn't exist, insert an empty node.
-		auto it = commands.find(child);
-		if (it == commands.end()) {
-			std::tie(it, std::ignore) =
-				commands.insert(std::pair{child,
-					node(std::string(child))});
-		}
-
-		if (path.empty())
-			return it->second;
-
-		return it->second.create_node(path);
-	}
-
-
-	void print_usage(std::string prefix) const {
-		if (handler != nullptr)
-			std::print("{}{}\n", prefix, name);
-
-		for (auto const &it : commands)
-			it.second.print_usage(prefix + name + " ");
+		return std::invoke(_func, argc, argv);
 	}
 
 private:
-	std::map<std::string_view, node> commands;
+	Fn _func;
 };
 
-// This may be called before main(), so catch any exceptions.
-template<typename Context>
-auto get_root_node() noexcept -> node<Context> & try {
-	static auto root_node = node<Context>("");
-	return root_node;
-} catch (std::exception const &exc) {
-	std::cerr << "ERROR: get_root_node: " << exc.what() << "\n";
-	std::exit(1);
-}
-
-export template<typename Context>
-auto dispatch_command(Context const &ctx, int argc, char **argv) -> int
-{
-	return get_root_node<Context>().invoke(ctx, argc, argv);
-}
-
-export template<typename Context>
-void print_usage(std::string_view prefix)
-{
-	get_root_node<Context>().print_usage(std::string(prefix));
-}
+// The public API.
+export auto dispatch_command(int argc, char **argv) -> int;
+export void print_usage(std::string_view prefix);
 
 } // namespace nihil
diff --git a/nihil/tests/command_map.cc b/nihil/tests/command_map.cc
index c8cd1a1..7d34b33 100644
--- a/nihil/tests/command_map.cc
+++ b/nihil/tests/command_map.cc
@@ -11,7 +11,7 @@ import nihil;
 namespace {
 
 auto cmd_sub1_called = false;
-auto cmd_sub1 = nihil::command<int>("cmd sub1", [](int, int, char **) -> int
+auto cmd_sub1 = nihil::command("cmd sub1", [](int, char **) -> int
 {
 	cmd_sub1_called = true;
 	return 0;
@@ -25,27 +25,6 @@ TEST_CASE("command_map: basic", "[command_map]")
 		"cmd", "sub1", nullptr
 	};
 	auto argv = const_cast<char **>(args.data());
-	nihil::dispatch_command(0, args.size() - 1, argv);
+	nihil::dispatch_command(args.size() - 1, argv);
 	REQUIRE(cmd_sub1_called == true);
 }
-
-TEST_CASE("command_map: no arguments", "[command_map]")
-{
-	auto args = std::vector<char const *>{
-		nullptr
-	};
-	auto argv = const_cast<char **>(args.data());
-	REQUIRE_THROWS_AS(nihil::dispatch_command(0, args.size() - 1, argv),
-			  nihil::usage_error);
-}
-
-TEST_CASE("command_map: unknown command", "[command_map]")
-{
-	auto args = std::vector<char const *>{
-		"cmd", "nonesuch", nullptr
-	};
-	auto argv = const_cast<char **>(args.data());
-
-	REQUIRE_THROWS_AS(nihil::dispatch_command(0, args.size() - 1, argv),
-			  nihil::usage_error);
-}