From e4c193005e81b6082a3698c000d60fcb516c131b Mon Sep 17 00:00:00 2001
From: Sam Harwell <sam@tunnelvisionlabs.com>
Date: Thu, 25 Sep 2014 21:02:13 -0500
Subject: [PATCH] Fix exponential complexity of recursive imports

---
 tool/src/org/antlr/v4/Tool.java         | 45 +++++++++++++++----------
 tool/src/org/antlr/v4/tool/Grammar.java | 26 +++++++++-----
 2 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/tool/src/org/antlr/v4/Tool.java b/tool/src/org/antlr/v4/Tool.java
index e181e04f3..26a42b2fb 100644
--- a/tool/src/org/antlr/v4/Tool.java
+++ b/tool/src/org/antlr/v4/Tool.java
@@ -608,6 +608,8 @@ public class Tool {
 		return g;
 	}
 
+	private final Map<String, Grammar> importedGrammars = new HashMap<String, Grammar>();
+
 	/**
 	 * Try current dir then dir of g then lib dir
 	 * @param g
@@ -615,27 +617,34 @@ public class Tool {
 	 */
 	public Grammar loadImportedGrammar(Grammar g, GrammarAST nameNode) throws IOException {
 		String name = nameNode.getText();
-		g.tool.log("grammar", "load " + name + " from " + g.fileName);
-		File importedFile = null;
-		for (String extension : ALL_GRAMMAR_EXTENSIONS) {
-			importedFile = getImportedGrammarFile(g, name + extension);
-			if (importedFile != null) {
-				break;
+		Grammar imported = importedGrammars.get(name);
+		if (imported == null) {
+			g.tool.log("grammar", "load " + name + " from " + g.fileName);
+			File importedFile = null;
+			for (String extension : ALL_GRAMMAR_EXTENSIONS) {
+				importedFile = getImportedGrammarFile(g, name + extension);
+				if (importedFile != null) {
+					break;
+				}
 			}
+
+			if ( importedFile==null ) {
+				errMgr.grammarError(ErrorType.CANNOT_FIND_IMPORTED_GRAMMAR, g.fileName, nameNode.getToken(), name);
+				return null;
+			}
+
+			String absolutePath = importedFile.getAbsolutePath();
+			ANTLRFileStream in = new ANTLRFileStream(absolutePath, grammarEncoding);
+			GrammarRootAST root = parse(g.fileName, in);
+			if (root == null) {
+				return null;
+			}
+
+			imported = createGrammar(root);
+			imported.fileName = absolutePath;
+			importedGrammars.put(root.getGrammarName(), imported);
 		}
 
-		if ( importedFile==null ) {
-			errMgr.grammarError(ErrorType.CANNOT_FIND_IMPORTED_GRAMMAR, g.fileName, nameNode.getToken(), name);
-			return null;
-		}
-
-		ANTLRFileStream in = new ANTLRFileStream(importedFile.getAbsolutePath(), grammarEncoding);
-		GrammarRootAST root = parse(g.fileName, in);
-		if ( root==null ) {
-			return null;
-		}
-		Grammar imported = createGrammar(root);
-		imported.fileName = importedFile.getAbsolutePath();
 		return imported;
 	}
 
diff --git a/tool/src/org/antlr/v4/tool/Grammar.java b/tool/src/org/antlr/v4/tool/Grammar.java
index 375394e61..e52d8d01f 100644
--- a/tool/src/org/antlr/v4/tool/Grammar.java
+++ b/tool/src/org/antlr/v4/tool/Grammar.java
@@ -514,16 +514,24 @@ public class Grammar implements AttributeResolver {
      *  The grammars are in import tree preorder.  Don't include ourselves
      *  in list as we're not a delegate of ourselves.
      */
-    public List<Grammar> getAllImportedGrammars() {
-        if ( importedGrammars==null ) return null;
-        List<Grammar> delegates = new ArrayList<Grammar>();
-		for (Grammar d : importedGrammars) {
-			delegates.add(d);
-			List<Grammar> ds = d.getAllImportedGrammars();
-			if (ds != null) delegates.addAll(ds);
+	public List<Grammar> getAllImportedGrammars() {
+		if (importedGrammars == null) {
+			return null;
 		}
-        return delegates;
-    }
+
+		LinkedHashMap<String, Grammar> delegates = new LinkedHashMap<String, Grammar>();
+		for (Grammar d : importedGrammars) {
+			delegates.put(d.fileName, d);
+			List<Grammar> ds = d.getAllImportedGrammars();
+			if (ds != null) {
+				for (Grammar imported : ds) {
+					delegates.put(imported.fileName, imported);
+				}
+			}
+		}
+
+		return new ArrayList<Grammar>(delegates.values());
+	}
 
     public List<Grammar> getImportedGrammars() { return importedGrammars; }