1 files changed, 88 insertions, 2 deletions
diff --git a/src/zencore/trace.cpp b/src/zencore/trace.cpp
index f7e4c4b68..ef7cbf596 100644
--- a/src/zencore/trace.cpp
+++ b/src/zencore/trace.cpp
@@ -4,10 +4,86 @@
 
 #	include <zencore/config.h>
 #	include <zencore/zencore.h>
+#	include <zencore/commandline.h>
+#	include <zencore/string.h>
+#	include <zencore/logging.h>
 
 #	define TRACE_IMPLEMENT 1
 #	include <zencore/trace.h>
 
+#	include <zencore/memory/memorytrace.h>
+
+namespace zen {
+
+void
+TraceConfigure()
+{
+	// Configure channels based on command line options
+
+	using namespace std::literals;
+
+	constexpr std::string_view TraceOption = "--trace="sv;
+
+	std::function<void(const std::string_view&)> ProcessChannelList;
+
+	auto ProcessTraceArg = [&](const std::string_view& Arg) {
+		if (Arg == "default"sv)
+		{
+			ProcessChannelList("cpu,log"sv);
+		}
+		else if (Arg == "memory"sv)
+		{
+			ProcessChannelList("memtag,memalloc,callstack,module"sv);
+		}
+		else if (Arg == "memory_light"sv)
+		{
+			ProcessChannelList("memtag,memalloc"sv);
+		}
+		else if (Arg == "memtag"sv)
+		{
+			// memtag actually traces to the memalloc channel
+			ProcessChannelList("memalloc"sv);
+		}
+		else
+		{
+			// Presume that the argument is a trace channel name
+
+			StringBuilder<128> AnsiChannel;
+			AnsiChannel << Arg;
+
+			const bool IsEnabled = trace::ToggleChannel(AnsiChannel.c_str(), true);
+
+			if (IsEnabled == false)
+			{
+				// Logging here could be iffy, but we might want some other feedback mechanism here
+				// to indicate to users that they're not getting what they might be expecting
+			}
+		}
+	};
+
+	ProcessChannelList = [&](const std::string_view& OptionArgs) { IterateCommaSeparatedValue(OptionArgs, ProcessTraceArg); };
+
+	bool TraceOptionPresent = false;
+
+	std::function<void(const std::string_view&)> ProcessArg = [&](const std::string_view& Arg) {
+		if (Arg.starts_with(TraceOption))
+		{
+			const std::string_view OptionArgs = Arg.substr(TraceOption.size());
+
+			TraceOptionPresent = true;
+
+			ProcessChannelList(OptionArgs);
+		}
+	};
+
+	IterateCommandlineArgs(ProcessArg);
+
+	if (!TraceOptionPresent)
+	{
+		ProcessTraceArg("default"sv);
+	}
+}
+
 void
 TraceInit(std::string_view ProgramName)
 {
@@ -38,6 +114,16 @@ TraceInit(std::string_view ProgramName)
 #	endif
 						   CommandLineString,
 						   ZEN_CFG_VERSION_BUILD_STRING);
+
+	atexit([] {
+#	if ZEN_WITH_MEMTRACK
+		zen::MemoryTrace_Shutdown();
+#	endif
+		trace::Update();
+		TraceShutdown();
+	});
+
+	TraceConfigure();
 }
 
 void
@@ -70,13 +156,11 @@ TraceStart(std::string_view ProgramName, const char* HostOrPath, TraceType Type)
 		case TraceType::None:
 			break;
 	}
-	trace::ToggleChannel("cpu", true);
 }
 
 bool
 TraceStop()
 {
-	trace::ToggleChannel("cpu", false);
 	if (trace::Stop())
 	{
 		return true;
@@ -84,4 +168,6 @@ TraceStop()
 	return false;
 }
 
+}  // namespace zen
+
 #endif	// ZEN_WITH_TRACE