Introduce the fund

src/collector/context.js

@@ -35,7 +35,7 @@ export default class Context {
             this._scope.addDefinition(name, type, declared);
         }
 
-        this._collector.types.push(type);
+        this._collector._fund.put(type);
 
         return type;
     }

src/collector/index.js

@@ -8,6 +8,7 @@ import traverse from './traverse';
 import globals from './globals';
 import definitionGroup from './definitions';
 import declarationGroup from './declarations';
+import Fund from '../fund';
 import Module from './module';
 import Scope from './scope';
 import Context from './context';
@@ -21,14 +22,14 @@ const VISITOR = Object.assign({}, definitionGroup, declarationGroup);
 export default class Collector {
     +root: string;
     +parser: Parser;
-    +types: Type[];
+    _fund: Fund;
     _modules: Map<string, Module>;
     _global: Scope;
 
     constructor(parser: Parser, root: string = '.') {
         this.root = root;
         this.parser = parser;
-        this.types = [];
+        this._fund = new Fund;
         this._modules = new Map;
         this._global = Scope.global(globals);
     }
@@ -63,6 +64,10 @@ export default class Collector {
         }
     }
 
+    finish(): Fund {
+        return this._fund;
+    }
+
     _freestyle(root: Node, scope: Scope, params: TemplateParam[]) {
         const ctx = new Context(this, scope, params);
 
@@ -146,37 +151,22 @@ export default class Collector {
                 return result.type;
 
             case 'special':
-                const resolve = id => this._findTypeById(id);
+            default:
+                const resolve = id => this._fund.take(id);
                 const type = result.call(params, resolve);
 
                 invariant(type);
 
                 return type;
         }
-
-        invariant(false);
     }
 
     _grabExports(module: Module) {
         for (const [scope, name] of module.exports()) {
-            this._query(scope, name, []);
-        }
-    }
-
-    _findTypeById(id: TypeId): Type {
-        // TODO: get rid of the linear search.
+            const type = this._query(scope, name, []);
 
-        for (const type of this.types) {
-            invariant(type.id);
-
-            const is = type.id.join('::') === id.join('::');
-
-            if (is) {
-                return type;
-            }
+            this._fund.put(type, true);
         }
-
-        invariant(false);
     }
 }
 

src/fund.js

@@ -0,0 +1,34 @@
+import {invariant} from './utils';
+import type {TypeId, Type} from './types';
+
+export default class Fund {
+    _tops: Type[];
+    _types: Map<string, Type>;
+
+    constructor() {
+        this._tops = [];
+        this._types = new Map;
+    }
+
+    put(type: Type, topLevel: boolean = false) {
+        invariant(type.id);
+
+        if (topLevel) {
+            this._tops.push(type);
+        }
+
+        this._types.set(JSON.stringify(type.id), type);
+    }
+
+    take(id: TypeId): Type {
+        const type = this._types.get(JSON.stringify(id));
+
+        invariant(type);
+
+        return type;
+    }
+
+    flatten(): Type[] {
+        return Array.from(this._types.values());
+    }
+}

src/index.js

@@ -11,5 +11,9 @@ export default function (path: string): {+types: Type[]} {
 
     collector.collect(path);
 
-    return collector;
+    const fund = collector.finish();
+
+    return {
+        types: fund.flatten(),
+    };
 }