Refactored tree-sitter playground.js

playground.js

// Dev switch for loading prev state or force canned eg.
let activateSaveState = true;
let showParseCount = true;

// Prelim sample input, drawn from the cli/src/tests/query_test.rs (as-is, excess space).
const eg = {
  lang: 'javascript',
  code: `
            class Person {
                // the constructor
                constructor(name) { this.name = name; }

                // the getter
                getFullName() { return this.name; }
            }
            `,
  query: `(class_declaration
                name: (identifier) @the-class-name
                (class_body
                    (method_definition
                        name: (property_identifier) @the-method-name
                        )))`
};
      
(async () => {
  const CAPTURE_REGEX = /@\s*([\w\._-]+)/g;
  const COLORS_BY_INDEX = [
    'blue',
    'chocolate',
    'darkblue',
    'darkcyan',
    'darkgreen',
    'darkred',
    'darkslategray',
    'dimgray',
    'green',
    'indigo',
    'navy',
    'red',
    'sienna',
  ];

//   const scriptURL = document.currentScript.getAttribute('src');

  const defaultCodeMirrorOpts = {
    lineNumbers: true,
    showCursorWhenSelecting: true
  }; 
    
  // State vars for Tree-sitter ops.
  const parse = {
    languageName: '', 
    language: null, // Tree-sitter Language obj for parser.setLanguage
    languages: {}, // language objs loaded, keyed by name
    parser: null,
    tree: null,
    query: null,
  
  };
  
  // State vars for ui event handling.
  const play = {
    //flags, counts, indices
    parseCount: 0,
    isRendering: 0,
    treeHighlight: -1,
    
    // UI components
    code: CodeMirror.fromTextArea(document.getElementById('code-input'), 
      defaultCodeMirrorOpts),
    query: CodeMirror.fromTextArea(document.getElementById('query-input'), 
      defaultCodeMirrorOpts),
    language: document.getElementById('language-select'),
    cluster: new Clusterize({
        rows: [],
        noDataText: null,
        contentElem: document.getElementById('output-container'),
        scrollElem: document.getElementById('output-container-scroll')
      }), 
    time: document.getElementById('update-time'),
    logEnable: document.getElementById('logging-checkbox'),
    queryEnable: document.getElementById('query-checkbox'),
  };

  const debRenderTree = debounce(renderTree, 50);
  const debSaveState = debounce(saveState, 2000);
  const debRunTreeQuery = debounce(runTreeQuery, 50);

  // CodeMirror ears.
  play.code.on('changes', handleCodeChange);
  play.code.on('viewportChange', debRunTreeQuery);
  play.code.on('cursorActivity', debounce(handleCodeCursorActivity, 150));
  play.query.on('changes', debounce(handleQueryChange, 150));

  // Dom ears.
  play.logEnable.addEventListener('change', handleLogEnableChange);
  play.queryEnable.addEventListener('change', handleQueryEnableChange);
  play.language.addEventListener('change', handleLanguageChange);

  play.cluster.content_elem.addEventListener('click', handleTreeClick);
  
  await TreeSitter.init();

  parse.parser = new TreeSitter();

  // Prime everything on empty doc.
  await handleLanguageChange();

  loadState();

  handleQueryEnableChange();

  document.getElementById('playground-container').style.visibility = 'visible';
  
  // Return language obj for name or null.
  async function loadLanguage(name) {
    if (!parse.languages[name]) {
      const url = `${LANGUAGE_BASE_URL}/tree-sitter-${name}.wasm`
      try {
        parse.languages[name] = await TreeSitter.Language.load(url);
      } catch (e) {
        console.error(e);
        return null;
      }
    }
    success = parse.parser.setLanguage(parse.languages[name]); 
    if (!success) {
      console.log('loadLanguage: setLanguage failed. name: ', name);
      return null;
    }
    parse.languageName = name;
    parse.language = parse.languages[name]
    parse.tree = null;
    return parse.language;
  }

  // When user chooses from language list
  async function handleLanguageChange() {
    const language = await loadLanguage(play.language.value);
    if (!language) {
      // Failed to find new lang, reset the view to the previous language.
      play.language.value = parse.languageName;
      return;
    }    
    handleCodeChange();
    handleQueryChange();
  }

  async function parseCode(newText, edits, timing=true) {
    let start = duration = null;
    if (timing) start = performance.now();
    if (edits) {
      for (const edit of edits) {
        parse.tree.edit(edit);
      }
    }
    const newTree = parse.parser.parse(newText, parse.tree);
    if (timing) duration = (performance.now() - start).toFixed(1);
    if (parse.tree) parse.tree.delete();
    parse.tree = newTree;  
    return duration; 
  }

  async function handleCodeChange(_, changes) {
    const newText = play.code.getValue() + '\n';
    const edits = parse.tree && changes && changes.map(gatherCodeEdits);

    let duration = await parseCode(newText, edits);
    play.parseCount++;
    if (duration) {
      const count = (showParseCount) ? ` (parse count: ${play.parseCount})` : ''
      play.time.innerText = `${duration} ms${count}`;
    }

    debRenderTree(); 
    debRunTreeQuery(); 
    debSaveState(); 
  }

  function composeDisplayName(cursor) {
    let name;
    if (cursor.nodeIsMissing) {
      name = `MISSING ${cursor.nodeType}`
    } else if (cursor.nodeIsNamed) {
      name = cursor.nodeType;
    }
    return name;
  }

  function composeRowStart(cursor, displayName, indentLevel) {
    // TSPosition {row, column}.
    const start = cursor.startPosition;
    const end = cursor.endPosition;
    const id = cursor.nodeId;
    let fieldName = cursor.currentFieldName();
    fieldName = (fieldName) ? fieldName + ': ' : '';

    let row = `<div>${'  '.repeat(indentLevel)}${fieldName}`;
    row += `<a class='plain' href="#" data-id=${id}`;
    row += ` data-range="${start.row},${start.column},${end.row},${end.column}">`;
    row += `${displayName}</a>`;
    row += ` [${start.row}, ${start.column}] - [${end.row}, ${end.column}]`;
    return row
  }
  
  function composeRowEnd(rows, row) {
    row += '</div>';
    rows.push(row);
    return '';
  }


  async function renderTree() {
    let rows = await composeTree() || [];
    play.cluster.update(rows);
    play.cluster.rows = rows;
    handleCodeCursorActivity();
  }

  // Walk the syntax tree and form html for each node, one node per row in the cluster.
  async function composeTree() {
    play.isRendering++;
    const cursor = parse.tree.walk();
    let currentRenderCount = play.parseCount;
    let row = '';
    let rows = [];
    let afterChildren = false;
    let indentLevel = 0;
    for (let i = 0;; i++) {
      if (i > 0 && i % 10000 === 0) {
        await new Promise(r => setTimeout(r, 0));
        if (play.parseCount !== currentRenderCount) {
          cursor.delete();
          play.isRendering--;
          return;
        }
      }

      let displayName = composeDisplayName(cursor);
      if (afterChildren) {
        if (cursor.gotoNextSibling()) {
          afterChildren = false;
        } else if (cursor.gotoParent()) {
          afterChildren = true;
          indentLevel--;
        } else {
          break;
        }
      } else {
        if (displayName) {
          if (row) row = composeRowEnd(rows, row);
          row = composeRowStart(cursor, displayName, indentLevel);
        }

        if (cursor.gotoFirstChild()) {
          afterChildren = false;
          indentLevel++;
        } else {
          afterChildren = true;
        }
      }
    }
    if (row) composeRowEnd(rows, row);

    cursor.delete();
    play.isRendering--;
    return rows;
  }

  // Mark (with color) the code per query result captures.
  function markCodePerQuery(captures) {
    play.code.operation(() => {
      play.code.getAllMarks().forEach(m => m.clear());

      let lastNodeId;
      for (const {name, node} of captures) {
        if (node.id === lastNodeId) continue;
        lastNodeId = node.id;

        // TSPosition {row, column} to CodeMirror {line, ch}
        const {startPosition, endPosition} = node;
        play.code.markText(
          {line: startPosition.row, ch: startPosition.column},
          {line: endPosition.row, ch: endPosition.column},
          {
            inclusiveLeft: true,
            inclusiveRight: true,
            css: `color: ${colorForCaptureName(name)}`
          }
        );
      }
    });
  }

  function runTreeQuery(_, startRow, endRow) {
    if (endRow == null) {
      const viewport = play.code.getViewport();
      startRow = viewport.from;
      endRow = viewport.to;
    }
    if (parse.tree && parse.query) {
      markCodePerQuery(parse.query.captures(parse.tree.rootNode,
        {row: startRow, column: 0}, {row: endRow, column: 0}));
    }
  }
  
  // Mark (with color) the query capture names.
  function markQuery() {
    play.query.operation(() => {
      play.query.getAllMarks().forEach(m => m.clear());

      let match;
      let row = 0;
      play.query.eachLine((line) => {
        while (match = CAPTURE_REGEX.exec(line.text)) {
          play.query.markText(
            {line: row, ch: match.index},
            {line: row, ch: match.index + match[0].length},
            {
              inclusiveLeft: true,
              inclusiveRight: true,
              css: `color: ${colorForCaptureName(match[1])}`
            }
          );
        }
        row++;
      });
    });
  }

  // Mark (with a CSS class) the invalid part of the query expression.
  function markQueryError(error, queryText) {
    play.query.operation(() => {
      play.query.getAllMarks().forEach(m => m.clear());

      // CodeMirror {line, ch}.
      const start = play.query.posFromIndex(error.index);
      const end = {line: start.line, ch: start.ch + (error.length || Infinity)};
      if (error.index === queryText.length) {
        if (start.ch > 0) {
          start.ch--;
        } else if (start.line > 0) {
          start.line--;
          start.ch = Infinity;
        }
      }
      play.query.markText(
        start,
        end,
        {
          className: 'query-error',
          inclusiveLeft: true,
          inclusiveRight: true,
          attributes: {title: error.message}
        }
      );
    });
  }

  function handleQueryChange() {
    if (parse.query) {
      parse.query.delete();
      parse.query.deleted = true;
      parse.query = null;
    }    

    if (!play.queryEnable.checked) return;

    const queryText = play.query.getValue();
    try {
      parse.query = parse.parser.getLanguage().query(queryText);
      markQuery();
    } catch (error) {
      markQueryError(play.query, error, queryText);
    }        

    runTreeQuery();
    saveQueryState();
  }

  function scrollTree() {
    const lineHeight = play.cluster.options.item_height;
    const scrollTop = play.cluster.scroll_elem.scrollTop;
    const containerHeight = play.cluster.scroll_elem.clientHeight;
    const offset = play.treeHighlight * lineHeight;
    if (scrollTop > offset - 20) {
      $(play.cluster.scroll_elem).animate({scrollTop: offset - 20}, 150);
    } else if (scrollTop < offset + lineHeight + 40 - containerHeight) {
      $(play.cluster.scroll_elem).animate(
        {scrollTop: offset - containerHeight + 40}, 150);
    }
  }

  // Ensure TSInputEdit will have start < end.
  function normalizeSelection(start, end) {
    // Tree-sitter {row, column}.
    if (start.row > end.row || (start.row === end.row && start.column > end.column)) {
      let swap = end;
      end = start;
      start = swap;
    }
    return [start, end];
  }
  
  function gatherCodeSelection() {
    const selection = play.code.getDoc().listSelections()[0];
    
    // CodeMirror {line, ch} to Tree-sitter {row, column}.
    let start = {row: selection.anchor.line, column: selection.anchor.ch};
    let end = {row: selection.head.line, column: selection.head.ch};
    return normalizeSelection(start, end);
  }

  function handleCodeCursorActivity() {
    // Ignore cursor if the tree is currently being rendered.
    if (play.isRendering) return;

    let [start, end] = gatherCodeSelection();
    const node = parse.tree.rootNode.namedDescendantForPosition(start, end);

    // Nothing more to do if tree pane is empty.
    if (!play.cluster.rows) return; 

    let rows = play.cluster.rows; 

    // Remove tree highlighting at current index, if any.
    let idx = play.treeHighlight;
    if (idx !== -1) {
      if (rows[idx]) rows[idx] = rows[idx].replace('highlighted', 'plain');
    }
  
    // Find new index and add tree highlighting, if any.
    idx = rows.findIndex(row => row.includes(`data-id=${node.id}`));
    if (idx !== -1) {
      if (rows[idx]) rows[idx] = rows[idx].replace('plain', 'highlighted');
    }
    play.treeHighlight = idx;
   
    play.cluster.update(rows);
    scrollTree();
  }
  
  function selectCode(startRow, startColumn, endRow, endColumn) {
    play.code.focus();
    play.code.setSelection(
      {line: startRow, ch: startColumn},
      {line: endRow, ch: endColumn}
    );
  }

  function handleTreeClick(event) {
    if (event.target.tagName === 'A') {
      event.preventDefault();
      const [startRow, startColumn, endRow, endColumn] = event
        .target
        .dataset
        .range
        .split(',')
        .map(n => parseInt(n));
      selectCode(startRow, startColumn, endRow, endColumn);
    }
  }

  function handleLogEnableChange() {
    if (play.logEnable.checked) {
      parse.parser.setLogger((message, lexing) => {
        (lexing) ? console.log("  ", message) : console.log(message);
      });
    } else {
      parse.parser.setLogger(null);
    }
  }

  function handleQueryEnableChange() {
    let elem = document.getElementById('query-container');
    if (play.queryEnable.checked) {
      elem.style.visibility = '';
      elem.style.position = '';
    } else {
      elem.style.visibility = 'hidden';
      elem.style.position = 'absolute';
    }
    handleQueryChange();
  }
  
  function gatherCodeEdits(change) {
    const oldLineCount = change.removed.length;
    const newLineCount = change.text.length;
    const lastLineLength = change.text[newLineCount - 1].length;

    // CodeMirror {line, ch} to Tree-sitter {row, column}.
    const startPosition = {row: change.from.line, column: change.from.ch};
    const oldEndPosition = {row: change.to.line, column: change.to.ch};
    const newEndPosition = {
      row: startPosition.row + newLineCount - 1,
      column: (newLineCount === 1)
        ? startPosition.column + lastLineLength
        : lastLineLength
    };

    const startIndex = play.code.indexFromPos(change.from);
    let newEndIndex = startIndex + newLineCount - 1;
    let oldEndIndex = startIndex + oldLineCount - 1;
    for (let i = 0; i < newLineCount; i++) newEndIndex += change.text[i].length;
    for (let i = 0; i < oldLineCount; i++) oldEndIndex += change.removed[i].length;

    // Tree-sitter InputEdit.
    return {
      startIndex, oldEndIndex, newEndIndex,
      startPosition, oldEndPosition, newEndPosition
    };
  }

  function colorForCaptureName(capture) {
    const id = parse.query.captureNames.indexOf(capture);
    return COLORS_BY_INDEX[id % COLORS_BY_INDEX.length];
  }

  function loadState() {
    if (!activateSaveState) {
      play.code.setValue(eg.code);
      play.query.setValue(eg.query);
      play.language.value = eg.lang; 
      play.queryEnable.checked = true;
      return
    }
    const language = localStorage.getItem("language");
    const sourceCode = localStorage.getItem("sourceCode");
    const query = localStorage.getItem("query");
    const queryEnabled = localStorage.getItem("queryEnabled");
    if (language != null && sourceCode != null && query != null) {
      play.code.setValue(sourceCode);
      play.query.setValue(query);
      play.language.value = language;
      play.queryEnable.checked = (queryEnabled === 'true');
    }
  }
  function saveState() {
    if (!activateSaveState) return;
    localStorage.setItem("language", play.language.value);
    localStorage.setItem("sourceCode", play.code.getValue());
    saveQueryState();
  }

  function saveQueryState() {
    if (!activateSaveState) return;
    localStorage.setItem("queryEnabled", play.queryEnable.checked);
    localStorage.setItem("query", play.query.getValue());
  }

  function debounce(func, wait, immediate) {
    var timeout;
    return function() {
      var context = this, args = arguments;
      var later = function() {
        timeout = null;
        if (!immediate) func.apply(context, args);
      };
      var callNow = immediate && !timeout;
      clearTimeout(timeout);
      timeout = setTimeout(later, wait);
      if (callNow) func.apply(context, args);
    };
  }
})();