/**
 * Copyright (c) 2016 Stefan Marr
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package tools.highlight;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import com.oracle.truffle.api.instrumentation.TruffleInstrument;
import com.oracle.truffle.api.instrumentation.TruffleInstrument.Registration;
import com.oracle.truffle.api.nodes.RootNode;
import com.oracle.truffle.api.source.SourceSection;

import tools.Tagging;


/**
 * <p>
 * The Highlight tool generates a JSON file with the source code that was
 * executed by the polyglot engine and uses {@link Tags} to provide information
 * for syntax highlighting based on source sections.
 *
 * <p>
 * Highlight expects that the root nodes that should be included are reported
 * to it using the {@link #reportParsedRootNode} function.
 * For elements that are not part of the AST, the {@link #reportNonAstSyntax}
 * function can be used.
 *
 * <p>
 * The output file can be configured by the Java property: hl.output
 */
@Registration(id = Highlight.ID)
public final class Highlight extends TruffleInstrument {

  public static final String ID = "highlight";

  private static final Map<SourceSection, Set<Class<? extends Tags>>> nonAstSourceSections = new HashMap<>();
  private static final Set<RootNode> rootNodes = new HashSet<>();

  /**
   * Report an element for highlighting that is not part of an AST,
   * and thus is not known to this tool from a root node
   * ({@link #reportParsedRootNode}).
   *
   * @param tag indicating the type of the element
   * @param source the source section describing the location of the element
   */
  // TODO: this is a bad hack. but, I don't know how to work around the polyglot engine otherwise
  //       normally, each polyglot engine should have a separate map
  public static void reportNonAstSyntax(final Class<? extends Tags> tag, final SourceSection source) {
    Set<Class<? extends Tags>> set = nonAstSourceSections.computeIfAbsent(source, s -> new HashSet<>(2));
    set.add(tag);
  }

  /**
   * Report a root node of which the AST nodes should be used for highlighting
   * code.
   *
   * @param rootNode created by the parser.
   *         While we would like to have all, there is not simply way to get them.
   */
  // TODO: this is a bad hack. but, I don't know how to work around the polyglot engine otherwise
  //       normally, each polyglot engine should have a separate map
  public static void reportParsedRootNode(final RootNode rootNode) {
    rootNodes.add(rootNode);
  }

  public Highlight() { }

  @Override
  protected void onCreate(final Env env) { }

  @Override
  protected void onDispose(final Env env) {
    // TODO: replace this by proper configuration from the environment
    String outputFile = System.getProperty("hl.output", "highlight.json");

    Map<SourceSection, Set<Class<? extends Tags>>> sourceSectionsAndTags = nonAstSourceSections;

    Tagging.collectSourceSectionsAndTags(rootNodes, sourceSectionsAndTags, env.getInstrumenter());
    JsonWriter.fileOut(outputFile, nonAstSourceSections);
  }
}

@@ -1229,24 +1229,84 @@
 .util.Set;%0A%0A
+import com.oracle.truffle.api.instrumentation.Instrumenter;%0A
 import com.o
@@ -1358,16 +1358,16 @@
 rument;%0A
-
 import c
@@ -2426,16 +2426,388 @@
 t%3C%3E();%0A%0A
+  private static Instrumenter instrumenter;%0A%0A  public static Map%3CSourceSection, Set%3CClass%3C? extends Tags%3E%3E%3E getSourceSections() %7B%0A    Map%3CSourceSection, Set%3CClass%3C? extends Tags%3E%3E%3E sourceSectionsAndTags = new HashMap%3C%3E(nonAstSourceSections);%0A    Tagging.collectSourceSectionsAndTags(rootNodes, sourceSectionsAndTags, instrumenter);%0A%0A    return sourceSectionsAndTags;%0A  %7D%0A%0A
   /**%0A  
@@ -4039,32 +4039,32 @@
  %7D%0A%0A  @Override%0A
-
   protected void
@@ -4089,16 +4089,60 @@
 v env) %7B
+%0A    instrumenter = env.getInstrumenter();%0A 
  %7D%0A%0A  @O

Add a getter for all source sections to Highlight tool