From 92202e1c8b4a7fee28e38fd964b6ce09321f38de Mon Sep 17 00:00:00 2001
From: Calvin Rose <calsrose@gmail.com>
Date: Sun, 21 Oct 2018 11:46:36 -0400
Subject: [PATCH] Add invert and simplify env-lookups to env-lookup

---
 src/core/core.janet       | 10 ++++++++++
 src/core/marsh.c          | 27 ++++-----------------------
 src/include/janet/janet.h |  3 +--
 test/suite1.janet         |  3 ++-
 4 files changed, 17 insertions(+), 26 deletions(-)

diff --git a/src/core/core.janet b/src/core/core.janet
index 3e6f1b3c..f629b9b9 100644
--- a/src/core/core.janet
+++ b/src/core/core.janet
@@ -710,6 +710,16 @@
      :tuple tuple.reverse
      :array array.reverse) t))
 
+(defn invert
+  "Returns a table of where the keys of an associative data structure
+are the values, and the values of the keys. If multiple keys have the same
+value, one key will be ignored."
+  [ds]
+  (def ret @{})
+  (loop [k :keys ds]
+    (put ret (get ds k) k))
+  ret)
+
 (defn zipcoll
   "Creates an table or tuple from two arrays/tuples. If a third argument of
   :struct is given result is struct else is table."
diff --git a/src/core/marsh.c b/src/core/marsh.c
index 64972ecc..8bb6735e 100644
--- a/src/core/marsh.c
+++ b/src/core/marsh.c
@@ -100,24 +100,8 @@ static Janet entry_getval(Janet env_entry) {
     }
 }
 
-/* Make a reverse lookup table for an environment (for marshaling) */
-JanetTable *janet_env_rreg(JanetTable *env) {
-    JanetTable *renv = janet_table(env->count);
-    while (env) {
-        for (int32_t i = 0; i < env->capacity; i++) {
-            if (janet_checktype(env->data[i].key, JANET_SYMBOL)) {
-                janet_table_put(renv,
-                        entry_getval(env->data[i].value),
-                        env->data[i].key);
-            }
-        }
-        env = env->proto;
-    }
-    return renv;
-}
-
 /* Make a forward lookup table from an environment (for unmarshaling) */
-JanetTable *janet_env_reg(JanetTable *env) {
+JanetTable *janet_env_lookup(JanetTable *env) {
     JanetTable *renv = janet_table(env->count);
     while (env) {
         for (int32_t i = 0; i < env->capacity; i++) {
@@ -1087,14 +1071,11 @@ int janet_unmarshal(
 
 /* C functions */
 
-static int cfun_env_lookups(JanetArgs args) {
+static int cfun_env_lookup(JanetArgs args) {
     JanetTable *env;
-    Janet tup[2];
     JANET_FIXARITY(args, 1);
     JANET_ARG_TABLE(env, args, 0);
-    tup[0] = janet_wrap_table(janet_env_rreg(env));
-    tup[1] = janet_wrap_table(janet_env_reg(env));
-    JANET_RETURN_TUPLE(args, janet_tuple_n(tup, 2));
+    JANET_RETURN_TABLE(args, janet_env_lookup(env));
 }
 
 static int cfun_marshal(JanetArgs args) {
@@ -1145,7 +1126,7 @@ static int cfun_unmarshal(JanetArgs args) {
 static const JanetReg cfuns[] = {
     {"marshal", cfun_marshal},
     {"unmarshal", cfun_unmarshal},
-    {"env-lookups", cfun_env_lookups},
+    {"env-lookup", cfun_env_lookup},
     {NULL, NULL}
 };
 
diff --git a/src/include/janet/janet.h b/src/include/janet/janet.h
index 451760a1..0860a20b 100644
--- a/src/include/janet/janet.h
+++ b/src/include/janet/janet.h
@@ -1037,8 +1037,7 @@ JANET_API int janet_unmarshal(
         Janet *out,
         JanetTable *reg,
         const uint8_t **next);
-JANET_API JanetTable *janet_env_rreg(JanetTable *env);
-JANET_API JanetTable *janet_env_reg(JanetTable *env);
+JANET_API JanetTable *janet_env_lookup(JanetTable *env);
 
 /* GC */
 JANET_API void janet_mark(Janet x);
diff --git a/test/suite1.janet b/test/suite1.janet
index 1f736ba6..dff646d7 100644
--- a/test/suite1.janet
+++ b/test/suite1.janet
@@ -141,7 +141,8 @@
 
 # Marshal
 
-(def [m-lookup um-lookup] (env-lookups _env))
+(def um-lookup (env-lookup _env))
+(def m-lookup (invert um-lookup))
 
 (defn testmarsh [x msg]
   (def marshx (marshal x m-lookup))