make more generic with grammar start rule and with file extension.

Author: parrt

java/src/org/antlr/codebuff/FeatureMetaDataTweaker.java

@@ -155,7 +155,7 @@ public static void main(String[] args) throws Exception {
 			corpusDir = "../samples/stringtemplate4/org/stringtemplate/v4/compiler/";
 			testFileDir = "../samples/stringtemplate4/org/stringtemplate/v4/compiler/";
 		}
-		Corpus corpus = Tool.train(corpusDir, JavaLexer.class, JavaParser.class, tabSize);
+		Corpus corpus = Tool.train(corpusDir, ".*\\.java", JavaLexer.class, JavaParser.class, "compilationUnit", tabSize);
 
 		List<String> allFiles = Tool.getFilenames(new File(testFileDir), ".*\\.java");
 		ArrayList<InputDocument> documents = (ArrayList<InputDocument>) Tool.load(allFiles, JavaLexer.class, tabSize);

java/src/org/antlr/codebuff/Optimizer.java

@@ -162,7 +162,7 @@ public static void main(String[] args) throws Exception {
 			corpusDir = "../samples/stringtemplate4/org/stringtemplate/v4/debug/";
 			testFileDir = "../samples/stringtemplate4/org/stringtemplate/v4/debug/";
 		}
-		Corpus corpus = Tool.train(corpusDir, JavaLexer.class, JavaParser.class, tabSize);
+		Corpus corpus = Tool.train(corpusDir, ".*\\.java", JavaLexer.class, JavaParser.class, "compilationUnit", tabSize);
 
 		List<String> allFiles = Tool.getFilenames(new File(testFileDir), ".*\\.java");
 		ArrayList<InputDocument> documents = (ArrayList<InputDocument>) Tool.load(allFiles, JavaLexer.class, tabSize);

java/src/org/antlr/codebuff/Tool.java

@@ -42,7 +42,7 @@ public static void main(String[] args)
 		String testFilename = args[2];
 		String output;
 		if ( language.equals("-java") ) {
-			Corpus corpus = train(corpusDir, JavaLexer.class, JavaParser.class, tabSize);
+			Corpus corpus = train(corpusDir, ".*\\.java", JavaLexer.class, JavaParser.class, "compilationUnit", tabSize);
 			InputDocument testDoc = load(testFilename, JavaLexer.class, tabSize);
 			Pair<String,List<TokenPositionAnalysis>> results = format(corpus, testDoc, tabSize);
 			output = results.a;
@@ -51,7 +51,7 @@ public static void main(String[] args)
 			controller.show();
 		}
 		else {
-			Corpus corpus = train(corpusDir, ANTLRv4Lexer.class, ANTLRv4Parser.class, tabSize);
+			Corpus corpus = train(corpusDir, ".*\\.g4", ANTLRv4Lexer.class, ANTLRv4Parser.class, "grammarSpec", tabSize);
 			InputDocument testDoc = load(testFilename, ANTLRv4Lexer.class, tabSize);
 			Pair<String,List<TokenPositionAnalysis>> results = format(corpus, testDoc, tabSize);
 			output = results.a;
@@ -89,18 +89,20 @@ public static Pair<String,List<TokenPositionAnalysis>> format(Corpus corpus, Inp
 	}
 
 	public static Corpus train(String rootDir,
+	                           String fileRegex,
 							   Class<? extends Lexer> lexerClass,
 							   Class<? extends Parser> parserClass,
+							   String startRuleName,
 							   int tabSize)
 		throws Exception
 	{
-		List<String> allFiles = getFilenames(new File(rootDir), ".*\\.java");
+		List<String> allFiles = getFilenames(new File(rootDir), fileRegex);
 		List<InputDocument> documents = load(allFiles, lexerClass, tabSize);
 
 		// Parse all documents into parse trees before training begins
 		for (InputDocument doc : documents) {
 			if ( showFileNames ) System.out.println(doc);
-			parse(doc, lexerClass, parserClass, "compilationUnit"); // TODO: make ruleName generic
+			parse(doc, lexerClass, parserClass, startRuleName);
 		}
 
 		// Walk all documents to compute matching token dependencies (we need this for feature computation)

java/test/org/antlr/codebuff/TestJavaStability.java

@@ -32,7 +32,7 @@ public TestJavaStability(String fileName) {
 
 	@Test
 	public void testStability() throws Exception {
-		Corpus corpus = Tool.train(fileName, JavaLexer.class, JavaParser.class, 4);
+		Corpus corpus = Tool.train(fileName, ".*\\.java", JavaLexer.class, JavaParser.class, "compilationUnit", 4);
 		InputDocument testDoc = Tool.load(fileName, JavaLexer.class, 4);
 		Pair<String,List<TokenPositionAnalysis>> results = Tool.format(corpus, testDoc, 4);
 		String output = results.a;