From 875e0f6766349f76baa7d38a51632b367e7d1d50 Mon Sep 17 00:00:00 2001
From: Taryn Hill <taryn@phrohdoh.com>
Date: Wed, 14 Jul 2021 17:16:06 -0500
Subject: [PATCH] feat: support clojure.core/get arity 3

https://clojure.github.io/clojure/clojure.core-api.html#clojure.core/get
> Returns the value mapped to key, not-found or nil if key not present.

    (get map key)           ;; supported prior to this commit
    (get map key not-found) ;; supported as of this commit
---
 src/persistent_list_map.rs       | 47 +++++++++++++++++++++++++++++++-
 src/protocols/ipersistent_map.rs |  8 ++++++
 src/rust_core/get.rs             | 14 ++++++----
 3 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/src/persistent_list_map.rs b/src/persistent_list_map.rs
index 486a192..6222490 100644
--- a/src/persistent_list_map.rs
+++ b/src/persistent_list_map.rs
@@ -14,7 +14,7 @@
 //! b => {:a 1 :b 3}
 
 use crate::maps::MapEntry;
-use crate::value::Value;
+use crate::value::{ToValue, Value};
 use crate::traits;
 
 use std::collections::HashMap;
@@ -101,6 +101,7 @@ macro_rules! merge {
 /// A PersistentListMap.
 pub trait IPersistentMap {
     fn get(&self, key: &Rc<Value>) -> Rc<Value>;
+    fn get_with_default(&self, key: &Rc<Value>, default: &Rc<Value>) -> Rc<Value>;
     fn assoc(&self, key: Rc<Value>, value: Rc<Value>) -> Self;
     fn contains_key(&self,key: &Rc<Value>) -> bool;
 }
@@ -117,6 +118,20 @@ impl IPersistentMap for PersistentListMap {
             PersistentListMap::Empty => Rc::new(Value::Nil),
         }
     }
+    fn get_with_default(&self, key: &Rc<Value>, default: &Rc<Value>) -> Rc<Value> {
+        match self {
+            PersistentListMap::Map(parent, entry) => {
+                if entry.key == *key {
+                    return Rc::clone(&entry.val);
+                }
+                match parent.get_with_default(key, default).as_ref() {
+                    Value::Nil => default.clone(),
+                    val => (*val).to_rc_value(),
+                }
+            }
+            PersistentListMap::Empty => default.clone(),
+        }
+    }
     fn assoc(&self, key: Rc<Value>, val: Rc<Value>) -> PersistentListMap {
         PersistentListMap::Map(Rc::new(self.clone()), MapEntry { key, val })
     }
@@ -146,6 +161,17 @@ impl IPersistentMap for Rc<PersistentListMap> {
             PersistentListMap::Empty => Rc::new(Value::Nil),
         }
     }
+    fn get_with_default(&self, key: &Rc<Value>, default: &Rc<Value>) -> Rc<Value> {
+        match &**self {
+            PersistentListMap::Map(parent, entry) => {
+                if entry.key == *key {
+                    return Rc::clone(&entry.val);
+                }
+                parent.get_with_default(key, default)
+            }
+            PersistentListMap::Empty => Rc::new(Value::Nil),
+        }
+    }
     fn assoc(&self, key: Rc<Value>, val: Rc<Value>) -> Rc<PersistentListMap> {
         Rc::new(PersistentListMap::Map(
             Rc::clone(self),
@@ -306,6 +332,25 @@ mod tests {
         println!("{}", map3);
         println!("{}", map4);
     }
+
+    #[test]
+    fn get_with_default() {
+        let map = persistent_list_map!(map_entry!("x", "v"));
+        let key = Keyword::intern("k").to_rc_value();
+        let default = Keyword::intern("not-found").to_rc_value();
+        let val: Rc<Value> = map.get_with_default(&key, &default);
+        assert_eq!(default, val);
+    }
+
+    #[test]
+    fn get_with_default_empty() {
+        let map = persistent_list_map!(/*empty*/);
+        let key = Keyword::intern("k").to_rc_value();
+        let default = Keyword::intern("not-found").to_rc_value();
+        let val: Rc<Value> = map.get_with_default(&key, &default);
+        assert_eq!(default, val);
+    }
+
     #[test]
     fn contains_key() {
         let map1 = persistent_list_map!{ "a" => 12, "b" => 13 };
diff --git a/src/protocols/ipersistent_map.rs b/src/protocols/ipersistent_map.rs
index 1f122f1..90c3da1 100644
--- a/src/protocols/ipersistent_map.rs
+++ b/src/protocols/ipersistent_map.rs
@@ -14,6 +14,14 @@ impl persistent_list_map::IPersistentMap for IPersistentMap {
             _ => panic!("Called Iterable iter on non-iterable"),
         }
     }
+    fn get_with_default(&self, key: &Rc<Value>, default: &Rc<Value>) -> Rc<Value> {
+        match &*self.value {
+            Value::PersistentListMap(plist_map) => {
+                plist_map.get_with_default(key, default)
+            },
+            _ => panic!("Called Iterable iter on non-iterable"),
+        }
+    }
     fn assoc(&self, key: Rc<Value>, value: Rc<Value>) -> IPersistentMap {
         match &*self.value {
             Value::PersistentListMap(plist_map) => {
diff --git a/src/rust_core/get.rs b/src/rust_core/get.rs
index 3dbe72c..e8c3777 100644
--- a/src/rust_core/get.rs
+++ b/src/rust_core/get.rs
@@ -1,3 +1,4 @@
+use crate::error_message;
 use crate::ifn::IFn;
 use crate::persistent_list_map::IPersistentMap;
 use crate::value::{ToValue, Value};
@@ -14,16 +15,17 @@ impl ToValue for GetFn {
 }
 impl IFn for GetFn {
     fn invoke(&self, args: Vec<Rc<Value>>) -> Value {
-        if args.len() != 2 {
-            return Value::Condition(format!(
-                "Wrong number of arguments given to function (Given: {}, Expected: 2)",
-                args.len()
-            ));
+        if args.len() != 2 && args.len() != 3 {
+            return error_message::wrong_varg_count(&[2, 3], args.len());
         }
 
         if let Value::PersistentListMap(pmap) = &*(args.get(0).unwrap().clone()) {
             let key = args.get(1).unwrap();
-            return pmap.get(key).to_value();
+            return if let Some(not_found) = args.get(2) {
+                pmap.get_with_default(key, not_found)
+            } else {
+                pmap.get(key)
+            }.to_value();
         }
         // @TODO add error in here with erkk's new error tools