From e3bac681ce25c5342087ecce1079e2c81a1acd8d Mon Sep 17 00:00:00 2001
From: Juraj Kirchheim <back2dos@googlemail.com>
Date: Fri, 24 Apr 2020 15:29:24 +0200
Subject: [PATCH] Allow for typemaps that reduce everything except null.

---
 src/tink/macro/TypeMap.hx | 33 ++++++++++++++++++++++++++-------
 tests/MyString.hx         |  1 +
 tests/TypeMapTest.hx      | 34 ++++++++++++++++++++++++++--------
 3 files changed, 53 insertions(+), 15 deletions(-)
 create mode 100644 tests/MyString.hx

diff --git a/src/tink/macro/TypeMap.hx b/src/tink/macro/TypeMap.hx
index 67e6996..1c44700 100644
--- a/src/tink/macro/TypeMap.hx
+++ b/src/tink/macro/TypeMap.hx
@@ -8,14 +8,33 @@ import haxe.macro.Type;
 using tink.MacroApi;
 
 class TypeMap<V> extends BalancedTree<Type, V> implements IMap<Type, V> {
-  var follow:Bool;
-  
-  public function new(?noFollow:Bool) {
-    this.follow = noFollow != true;
+  var normalize:Type->Type;
+
+  public function new(?normalize:Type->Type) {
+    this.normalize = switch normalize {
+      case null: function (t) return t.reduce();
+      case fn: fn;
+    }
     super();
   }
-  
+
   override function compare(k1:Type, k2:Type):Int
-    return k1.compare(k2, follow);
-  
+    return normalize(k1).compare(normalize(k2), false);
+
+  static public function noFollow(t:Type)
+    return t;
+
+  static public function keepNull(t:Type):Type
+    return switch t {
+      case TAbstract(_.get() => { name: 'Null', pack: [] }, [t])
+        #if !haxe4 | TType(_.get() => { name: 'Null', pack: []}, [t]) #end
+        :
+        var ct = keepNull(t).toComplex({ direct: true });
+        (macro : Null<$ct>).toType().sure();
+
+      case TLazy(f): keepNull(f());
+      case TType(_, _): keepNull(Context.follow(t, true));
+      default: t;
+    }
+
 }
\ No newline at end of file
diff --git a/tests/MyString.hx b/tests/MyString.hx
new file mode 100644
index 0000000..3e0cd75
--- /dev/null
+++ b/tests/MyString.hx
@@ -0,0 +1 @@
+typedef MyString = String;
\ No newline at end of file
diff --git a/tests/TypeMapTest.hx b/tests/TypeMapTest.hx
index 14166d6..201dd9c 100644
--- a/tests/TypeMapTest.hx
+++ b/tests/TypeMapTest.hx
@@ -1,17 +1,17 @@
 package;
 
 import haxe.unit.TestCase;
+import haxe.macro.Expr;
 
-import tink.macro.TypeMap;
-using haxe.macro.Context;
+using tink.MacroApi;
 
 class TypeMapTest extends TestCase {
 
   function testMap() {
     var t = new TypeMap();
-    var t1 = (macro [{ foo: [{ bar: '5' }]}]).typeof();
-    var t2 = (macro [{ foo: [{ bar: 5 }]}]).typeof();
-    
+    var t1 = (macro [{ foo: [{ bar: '5' }]}]).typeof().sure();
+    var t2 = (macro [{ foo: [{ bar: 5 }]}]).typeof().sure();
+
     t.set(t1, 0);
     assertEquals(Lambda.count(t), 1);
     t.set(t2, 1);
@@ -20,11 +20,29 @@ class TypeMapTest extends TestCase {
     assertEquals(Lambda.count(t), 2);
     t.set(t2, 3);
     assertEquals(Lambda.count(t), 2);
-    
+
     assertEquals(t.get(t1), 2);
     assertEquals(t.get(t2), 3);
-    
+
     assertTrue(true);
   }
-  
+
+  function testNormalization() {
+    for (settings in [
+      { count: 4, map: new TypeMap(TypeMap.noFollow)},
+      { count: 2, map: new TypeMap(TypeMap.keepNull)},
+      { count: 1, map: new TypeMap()},
+    ]) {
+
+      function set(ct:ComplexType)
+        settings.map.set(ct.toType().sure(), ct);
+
+      set(macro : String);
+      set(macro : MyString);
+      set(macro : Null<String>);
+      set(macro : Null<MyString>);
+
+      assertEquals(settings.count, Lambda.count(settings.map));
+    }
+  }
 }
\ No newline at end of file