From f49d00ad35e96f155bf035bb0fd329c24d78ce34 Mon Sep 17 00:00:00 2001
From: pancake <pancake@nopcode.org>
Date: Tue, 25 Apr 2023 01:26:02 +0200
Subject: [PATCH] Improve automatch plugin with encoder support ##arch

---
 libr/arch/arch.c      | 37 +++++++++++++++++++++++++++++++------
 libr/include/r_arch.h |  4 +++-
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/libr/arch/arch.c b/libr/arch/arch.c
index fb9be6d72e..d96fd14d7d 100644
--- a/libr/arch/arch.c
+++ b/libr/arch/arch.c
@@ -51,12 +51,15 @@ static ut32 _rate_compat(RArchPlugin *p, RArchConfig *cfg, const char *name) {
 	return score;
 }
 
-static RArchPlugin *find_bestmatch(RArch *arch, RArchConfig *cfg, const char *name) {
+static RArchPlugin *find_bestmatch(RArch *arch, RArchConfig *cfg, const char *name, bool enc) {
 	ut8 best_score = 0;
 	RArchPlugin *ap = NULL;
 	RListIter *iter;
 	RArchPlugin *p;
 	r_list_foreach (arch->plugins, iter, p) {
+		if (enc && !p->encode) {
+			continue;
+		}
 		const ut32 score = _rate_compat (p, cfg, name);
 		if (score > 0 && score > best_score) {
 			best_score = score;
@@ -83,7 +86,7 @@ R_API bool r_arch_use(RArch *arch, RArchConfig *config, const char *name) {
 		return true;
 	}
 #endif
-	RArchPlugin *ap = find_bestmatch (arch, config, name);
+	RArchPlugin *ap = find_bestmatch (arch, config, name, false);
 	if (!ap) {
 		r_unref (arch->session);
 		arch->session = NULL;
@@ -91,6 +94,22 @@ R_API bool r_arch_use(RArch *arch, RArchConfig *config, const char *name) {
 	}
 	r_unref (arch->session);
 	arch->session = r_arch_session (arch, config, ap);
+	if (arch->session) {
+		RArchPluginEncodeCallback encode = arch->session->plugin->encode;
+		if (!encode) {
+#if R2_590
+			RArchPlugin *ap = find_bestmatch (arch, config, name, true);
+			if (ap) {
+				RArchSession *es = r_arch_session (arch, config, ap);
+				if (es && es->plugin == arch->session->plugin) {
+					r_unref (es);
+				} else if (es) {
+					arch->session->encoder = es;
+				}
+			}
+#endif
+		}
+	}
 #if 0
 	RArchConfig *oconfig = arch->cfg;
 	r_unref (arch->cfg);
@@ -193,7 +212,7 @@ R_API bool r_arch_add(RArch *a, RArchPlugin *ap) {
 
 R_API bool r_arch_del(RArch *arch, const char *name) {
 	r_return_val_if_fail (arch && arch->plugins && name, false);
-	RArchPlugin *ap = find_bestmatch (arch, NULL, name);
+	RArchPlugin *ap = find_bestmatch (arch, NULL, name, false);
 #if 0
 	if (arch->current && !strcmp (arch->current->p->name, name)) {
 		arch->current = NULL;
@@ -219,9 +238,15 @@ R_API int r_arch_info(RArch *a, int query) {
 }
 
 R_API bool r_arch_encode(RArch *a, RAnalOp *op, RArchEncodeMask mask) {
-	// XXX should be unused
-	RArchPluginEncodeCallback encode = R_UNWRAP4 (a, session, plugin, encode);
-	return encode? encode (a->session, op, mask): false;
+	RArchSession *session = a->session;
+	RArchPluginEncodeCallback encode = R_UNWRAP3 (session, plugin, encode);
+#if R2_590
+	if (!encode && session->encoder) {
+		session = session->encoder;
+		encode = R_UNWRAP3 (session, plugin, encode);
+	}
+#endif
+	return encode? encode (session, op, mask): false;
 }
 
 R_API bool r_arch_decode(RArch *a, RAnalOp *op, RArchDecodeMask mask) {
diff --git a/libr/include/r_arch.h b/libr/include/r_arch.h
index 7389b486f0..01937c743a 100644
--- a/libr/include/r_arch.h
+++ b/libr/include/r_arch.h
@@ -148,9 +148,11 @@ typedef struct r_arch_t {
 typedef struct r_arch_session_t {
 #if R2_590
 	char *name; // used by .use to chk if it was set already
+	// TODO: name it "peer" instead of encoder. so the encoder can back reference the decoder
+	struct r_arch_session_t *encoder; // used for encoding when plugin->encode is not set
 #endif
 	struct r_arch_t *arch;
-	struct r_arch_plugin_t *plugin;
+	struct r_arch_plugin_t *plugin; // used for decoding
 	RArchConfig *config; // TODO remove arch->config!
 	void *data;
 	void *user;