Support union types when adding argument-function-call effects

Issue #734
ternjs · Apr 8, 2016 · d039c3c · d039c3c
1 parent 35ea365
commit d039c3c
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/lib/def.js b/lib/def.js
@@ -244,15 +244,24 @@
     }
   };
 
+  function addArgCallEffects(type) {
+    if (type instanceof infer.Fn && type.args) for (var i = 0; i < type.args.length; ++i) {
+      let arg = type.args[i]
+      if (arg instanceof infer.Fn && arg.args && arg.args.length) addArgCallEffect(type, i)
+    }
+  }
+
+  function addArgCallEffect(type, argNum) {
+    addEffect(type, function(_self, args) {
+      if (args[argNum]) args[argNum].propagate(
+        new infer.IsCallee(infer.cx().topScope, type.args[argNum].args, null, infer.ANull))
+    })
+  }
+
   function parseType(spec, name, base, forceNew) {
     var type = new TypeParser(spec, null, base, forceNew).parseType(false, name, true);
-    if (/^fn\(/.test(spec)) for (var i = 0; i < type.args.length; ++i) (function(i) {
-      var arg = type.args[i];
-      if (arg instanceof infer.Fn && arg.args && arg.args.length) addEffect(type, function(_self, fArgs) {
-        var fArg = fArgs[i];
-        if (fArg) fArg.propagate(new infer.IsCallee(infer.cx().topScope, arg.args, null, infer.ANull));
-      });
-    })(i);
+    if (type instanceof infer.AVal) type.types.forEach(addArgCallEffects)
+    else addArgCallEffects(type)
     return type;
   }