Merge pull request #775 from quantified-uncertainty/semantic-highlight

Semantic highlighting

packages/squiggle-lang/.gitignore

@@ -22,3 +22,4 @@ _coverage
 coverage
 .nyc_output/
 src/rescript/Reducer/Reducer_Peggy/Reducer_Peggy_GeneratedParser.js
+src/rescript/Reducer/Reducer_Peggy/helpers.js

packages/squiggle-lang/package.json

@@ -7,7 +7,8 @@
     "peggy": "peggy --cache",
     "rescript": "rescript",
     "build": "yarn build:peggy && yarn build:rescript && yarn build:typescript",
-    "build:peggy": "find . -type f -name *.peggy -exec yarn peggy {} \\;",
+    "build:peggy:helpers": "tsc --module commonjs --outDir src/rescript/Reducer/Reducer_Peggy/ src/rescript/Reducer/Reducer_Peggy/helpers.ts",
+    "build:peggy": "yarn build:peggy:helpers && find . -type f -name *.peggy -exec yarn peggy {} \\;",
     "build:rescript": "rescript build -with-deps",
     "build:typescript": "tsc",
     "bundle": "webpack",

packages/squiggle-lang/src/js/index.ts

@@ -10,7 +10,6 @@ import {
   evaluatePartialUsingExternalBindings,
   evaluateUsingOptions,
   foreignFunctionInterface,
-  parse as parseRescript,
 } from "../rescript/TypescriptInterface.gen";
 export {
   makeSampleSetDist,
@@ -32,12 +31,14 @@ import {
   convertRawToTypescript,
   lambdaValue,
 } from "./rescript_interop";
-import { Ok, result, resultMap, tag, tagged } from "./types";
+import { result, resultMap, tag, tagged } from "./types";
 import { Distribution, shape } from "./distribution";
 
 export { Distribution, resultMap, defaultEnvironment };
 export type { result, shape, environment, lambdaValue, squiggleExpression };
 
+export { parse } from "./parse";
+
 export let defaultSamplingInputs: environment = {
   sampleCount: 10000,
   xyPointLength: 10000,
@@ -59,23 +60,6 @@ export function run(
   return resultMap(res, (x) => createTsExport(x, e));
 }
 
-export function parse(
-  squiggleString: string
-): result<null, Extract<errorValue, { tag: "RESyntaxError" }>> {
-  const maybeExpression = parseRescript(squiggleString);
-  if (maybeExpression.tag === "Ok") {
-    return Ok(null); // TODO - return AST
-  } else {
-    if (
-      typeof maybeExpression.value !== "object" ||
-      maybeExpression.value.tag !== "RESyntaxError"
-    ) {
-      throw new Error("Expected syntax error");
-    }
-    return { tag: "Error", value: maybeExpression.value };
-  }
-}
-
 // Run Partial. A partial is a block of code that doesn't return a value
 export function runPartial(
   squiggleString: string,

packages/squiggle-lang/src/js/parse.ts

@@ -0,0 +1,23 @@
+import {
+  errorValue,
+  parse as parseRescript,
+} from "../rescript/TypescriptInterface.gen";
+import { result } from "./types";
+import { AnyPeggyNode } from "../rescript/Reducer/Reducer_Peggy/helpers";
+
+export function parse(
+  squiggleString: string
+): result<AnyPeggyNode, Extract<errorValue, { tag: "RESyntaxError" }>> {
+  const maybeExpression = parseRescript(squiggleString);
+  if (maybeExpression.tag === "Ok") {
+    return { tag: "Ok", value: maybeExpression.value as AnyPeggyNode };
+  } else {
+    if (
+      typeof maybeExpression.value !== "object" ||
+      maybeExpression.value.tag !== "RESyntaxError"
+    ) {
+      throw new Error("Expected syntax error");
+    }
+    return { tag: "Error", value: maybeExpression.value };
+  }
+}

packages/squiggle-lang/src/rescript/Reducer/Reducer_Peggy/Reducer_Peggy_GeneratedParser.peggy

@@ -1,74 +1,7 @@
 // Try in https://peggyjs.org/online
 
 {{
-  var toFunction = {
+  const h = require('./helpers');
-    '-': 'subtract',
-    '->': 'pipe',
-    '!=': 'unequal',
-    '.-': 'dotSubtract',
-    '.*': 'dotMultiply',
-    './': 'dotDivide',
-    '.^': 'dotPow',
-    '.+': 'dotAdd',
-    '*': 'multiply',
-    '/': 'divide',
-    '&&': 'and',
-    '^': 'pow', // or xor
-    '+': 'add',
-    '<': 'smaller',
-    '<=': 'smallerEq',
-    '==': 'equal',
-    '>': 'larger',
-    '>=': 'largerEq',
-    '||': 'or',
-    'to': 'credibleIntervalToDistribution',
-  } 
-
-  var unaryToFunction = {
-    '-': 'unaryMinus', 
-    '!': 'not',
-    '.-': 'unaryDotMinus',
-  }
-
-  var postOperatorToFunction = {
-    '.': '$_atIndex_$',
-    '()': '$$_applyAll_$$',
-    '[]': '$_atIndex_$',
-  }
-
-  function makeFunctionCall(fn, args) {
-    if (fn === '$$_applyAll_$$') {
-      // Any list of values is applied from left to right anyway. 
-      // Like in Haskell and Lisp.
-      // So we remove the redundant $$_applyAll_$$.
-      if (args[0].type === "Identifier") {args[0].type = "CallIdentifier"}
-      return nodeExpression(args) 
-    } else {
-      return nodeExpression([nodeCallIndentifier(fn), ...args]) 
-    }
-  }
-
-  function apply(fn, arg) { return makeFunctionCall(fn, [arg]); }
-  function constructArray(elems)  { return apply('$_constructArray_$', nodeExpression(elems)); }
-  function constructRecord(elems) { return apply('$_constructRecord_$', nodeExpression(elems)); }
-
-  function nodeBlock(statements) {return{type: 'Block', statements: statements}}
-  function nodeBoolean(value) {return {type: 'Boolean', value: value}}
-  function nodeCallIndentifier(value) {return {type: 'CallIdentifier', value: value}}
-  function nodeExpression(args) {return {type: 'Expression', nodes: args}}
-  function nodeFloat(value) {return {type: 'Float', value: value}}
-  function nodeIdentifier(value) {return {type: 'Identifier', value: value}}
-  function nodeInteger(value) {return {type: 'Integer', value: value}}
-  function nodeKeyValue(key, value) { 
-    if (key.type === 'Identifier') {key.type = 'String'}
-    return {type: 'KeyValue', key: key, value: value}}  
-  function nodeLambda(args, body) {return {type: 'Lambda', args: args, body: body}}
-  function nodeLetStatment(variable, value) {return {type: 'LetStatement', variable: variable, value: value}}
-  function nodeModuleIdentifier(value) {return {type: 'ModuleIdentifier', value: value}}
-  function nodeString(value) {return {type: 'String', value: value}}
-  function nodeTernary(condition, trueExpression, falseExpression) {return {type: 'Ternary', condition: condition, trueExpression: trueExpression, falseExpression: falseExpression}}
-  
-  function nodeTypeIdentifier(typeValue) {return {type: 'TypeIdentifier', value: typeValue}}
 }}
 
 start
@@ -80,21 +13,21 @@ zeroOMoreArgumentsBlockOrExpression = innerBlockOrExpression / lambda
 outerBlock 
 	= statements:array_statements  finalExpression: (statementSeparator @expression)?
     { if (finalExpression != null) { statements.push(finalExpression) }
-        return nodeBlock(statements) }
+        return h.nodeBlock(statements) }
   / finalExpression: expression
-    { return nodeBlock([finalExpression])}
+    { return h.nodeBlock([finalExpression])}
     
 innerBlockOrExpression  
   = quotedInnerBlock
   / finalExpression: expression
-    { return nodeBlock([finalExpression])}
+    { return h.nodeBlock([finalExpression])}
 
 quotedInnerBlock  
   = '{' _nl statements:array_statements  finalExpression: (statementSeparator @expression)  _nl '}'
 	  { statements.push(finalExpression) 
-    	return nodeBlock(statements) }
+    	return h.nodeBlock(statements) }
   / '{' _nl finalExpression: expression  _nl '}'
-	  { return nodeBlock([finalExpression]) }
+	  { return h.nodeBlock([finalExpression]) }
 
 array_statements
   = head:statement tail:(statementSeparator @array_statements )
@@ -109,12 +42,12 @@ statement
 
 letStatement 
   = variable:identifier _ assignmentOp _nl value:zeroOMoreArgumentsBlockOrExpression
-  { return nodeLetStatment(variable, value) }
+  { return h.nodeLetStatement(variable, value) }
 
 defunStatement
   = variable:identifier '(' _nl args:array_parameters _nl ')' _ assignmentOp _nl body:innerBlockOrExpression
-    { var value = nodeLambda(args, body)
+    { var value = h.nodeLambda(args, body)
-      return nodeLetStatment(variable, value) }
+      return h.nodeLetStatement(variable, value) }
 
   assignmentOp "assignment" = '='
 
@@ -128,16 +61,16 @@ ifthenelse
   = 'if' __nl condition:logicalAdditive 
   	__nl 'then' __nl trueExpression:innerBlockOrExpression 
     __nl 'else' __nl falseExpression:(ifthenelse/innerBlockOrExpression)
-    { return nodeTernary(condition, trueExpression, falseExpression) }
+    { return h.nodeTernary(condition, trueExpression, falseExpression) }
   
 ternary 
   = condition:logicalAdditive _ '?' _nl trueExpression:logicalAdditive _ ':' _nl falseExpression:(ternary/logicalAdditive)
-    { return nodeTernary(condition, trueExpression, falseExpression) }
+    { return h.nodeTernary(condition, trueExpression, falseExpression) }
 
 logicalAdditive
   = head:logicalMultiplicative tail:(_ operator:logicalAdditiveOp _nl arg:logicalMultiplicative {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
   logicalAdditiveOp "operator" = '||'
@@ -146,21 +79,21 @@ logicalAdditive
 logicalMultiplicative
   = head:equality tail:(_ operator:logicalMultiplicativeOp _nl arg:equality {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
   logicalMultiplicativeOp "operator" = '&&'
 
 equality
   = left:relational _ operator:equalityOp _nl right:relational 
-  { return makeFunctionCall(toFunction[operator], [left, right])}
+  { return h.makeFunctionCall(h.toFunction[operator], [left, right])}
   / relational  
  
  equalityOp "operator" = '=='/'!='
 
 relational
   = left:additive _ operator:relationalOp _nl right:additive 
-  { return makeFunctionCall(toFunction[operator], [left, right])}
+  { return h.makeFunctionCall(h.toFunction[operator], [left, right])}
   / additive
 
   relationalOp "operator" = '<='/'<'/'>='/'>'
@@ -168,7 +101,7 @@ relational
 additive
   = head:multiplicative tail:(_ operator:additiveOp _nl arg:multiplicative {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
   additiveOp "operator" = '+' / '-' / '.+' / '.-'
@@ -176,7 +109,7 @@ additive
 multiplicative
   = head:power tail:(_ operator:multiplicativeOp _nl arg:power {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
   multiplicativeOp "operator" = '*' / '/' / '.*' / './'
@@ -184,7 +117,7 @@ multiplicative
 power
   = head:credibleInterval tail:(_ operator:powerOp _nl arg:credibleInterval {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
   powerOp "operator" = '^' / '.^'
@@ -192,7 +125,7 @@ power
 credibleInterval
   = head:chainFunctionCall tail:(__ operator:credibleIntervalOp __nl arg:chainFunctionCall {return {operator: operator, right: arg}})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(toFunction[element.operator], [result, element.right])
+      return h.makeFunctionCall(h.toFunction[element.operator], [result, element.right])
     }, head)}
 
     credibleIntervalOp "operator" = 'to'
@@ -200,7 +133,7 @@ credibleInterval
 chainFunctionCall
   = head:unary tail:(_ ('->'/'|>') _nl chained:chainedFunction {return chained})* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(element.fnName, [result, ...element.args])
+      return h.makeFunctionCall(element.fnName, [result, ...element.args])
     }, head)}
 
   chainedFunction
@@ -215,7 +148,7 @@ chainFunctionCall
 
 unary
   = unaryOperator:unaryOperator _nl right:(unary/postOperator)
-  { return apply(unaryToFunction[unaryOperator], right)}
+  { return h.apply(h.unaryToFunction[unaryOperator], right)}
   / postOperator
   
   unaryOperator "unary operator"
@@ -230,12 +163,12 @@ indexedValue
 collectionElement
   = head:atom &('['/'('/'.')
     tail:(
-      _ '[' _nl arg:expression  _nl ']' {return {fn: postOperatorToFunction['[]'], args: [arg]}}
+      _ '[' _nl arg:expression  _nl ']' {return {fn: h.postOperatorToFunction['[]'], args: [arg]}}
-    / _ '(' _nl args:array_functionArguments  _nl ')' {return {fn: postOperatorToFunction['()'], args: args}}
+    / _ '(' _nl args:array_functionArguments  _nl ')' {return {fn: h.postOperatorToFunction['()'], args: args}}
-    / '.' arg:$dollarIdentifier {return {fn: postOperatorToFunction['[]'], args: [nodeString(arg)]}}
+    / '.' arg:$dollarIdentifier {return {fn: h.postOperatorToFunction['[]'], args: [h.nodeString(arg)]}}
     )* 
   { return tail.reduce(function(result, element) {
-      return makeFunctionCall(element.fn, [result, ...element.args])
+      return h.makeFunctionCall(element.fn, [result, ...element.args])
     }, head)}
 
     array_functionArguments 
@@ -261,49 +194,49 @@ dollarIdentifierWithModule 'identifier'
     final:$dollarIdentifier
   { tail.push(final);
     return tail.reduce(function(result, element) {
-      return makeFunctionCall(postOperatorToFunction['[]'], [result, nodeString(element)])
+      return h.makeFunctionCall(h.postOperatorToFunction['[]'], [result, h.nodeString(element)])
     }, head)}
 
 identifier 'identifier'
-  = ([_a-z]+[_a-z0-9]i*) {return nodeIdentifier(text())} 
+  = ([_a-z]+[_a-z0-9]i*) {return h.nodeIdentifier(text(), location())} 
 
 unitIdentifier 'identifier'
-  = ([_a-zA-Z]+[_a-z0-9]i*) {return nodeIdentifier(text())} 
+  = ([_a-zA-Z]+[_a-z0-9]i*) {return h.nodeIdentifier(text(), location())} 
 
 dollarIdentifier '$identifier'
-  = ([\$_a-z]+[\$_a-z0-9]i*) {return nodeIdentifier(text())} 
+  = ([\$_a-z]+[\$_a-z0-9]i*) {return h.nodeIdentifier(text(), location())} 
 
 moduleIdentifier 'identifier'
-  = ([A-Z]+[_a-z0-9]i*) {return nodeModuleIdentifier(text())} 
+  = ([A-Z]+[_a-z0-9]i*) {return h.nodeModuleIdentifier(text())} 
 
 
 string 'string'
-  = characters:("'" @([^'])* "'") {return nodeString(characters.join(''))} 
+  = characters:("'" @([^'])* "'") {return h.nodeString(characters.join(''))} 
-  / characters:('"' @([^"])* '"') {return nodeString(characters.join(''))}
+  / characters:('"' @([^"])* '"') {return h.nodeString(characters.join(''))}
 
 number = number:(float / integer) unit:unitIdentifier?
   { 
     if (unit === null)
       { return number }
     else
-      { return apply('fromUnit_'+unit.value, number) 
+      { return h.apply('fromUnit_'+unit.value, number) 
       }
   }
 
 integer 'integer'
   = d+ !"\." ![e]i
-  { return nodeInteger(parseInt(text()))} 
+  { return h.nodeInteger(parseInt(text()))} 
   
 float 'float'
   = $(((d+ "\." d*) / ("\." d+)) floatExponent? / d+ floatExponent)
-  { return nodeFloat(parseFloat(text()))} 
+  { return h.nodeFloat(parseFloat(text()))} 
 
 	floatExponent = [e]i '-'? d+
 	d = [0-9]
     
 boolean 'boolean'
   = ('true'/'false') 
-  { return nodeBoolean(text() === 'true')}   
+  { return h.nodeBoolean(text() === 'true')}   
 
 valueConstructor
   = recordConstructor
@@ -314,15 +247,15 @@ valueConstructor
 lambda  
   = '{' _nl '|' _nl args:array_parameters _nl '|' _nl statements:array_statements  finalExpression: (statementSeparator @expression)  _nl '}'
 	  { statements.push(finalExpression) 
-    	return nodeLambda(args, nodeBlock(statements)) }
+    	return h.nodeLambda(args, h.nodeBlock(statements)) }
   / '{' _nl '|' _nl args:array_parameters _nl '|' _nl finalExpression: expression  _nl '}'
-	  { return nodeLambda(args, nodeBlock([finalExpression])) }
+	  { return h.nodeLambda(args, h.nodeBlock([finalExpression])) }
 
 arrayConstructor 'array'
   = '[' _nl ']'
-    { return constructArray([]); }
+    { return h.constructArray([]); }
   / '[' _nl args:array_elements _nl ']' 
-    { return constructArray(args); }
+    { return h.constructArray(args); }
 
   array_elements 
     = head:expression tail:(_ ',' _nl @expression)* 
@@ -330,7 +263,7 @@ arrayConstructor 'array'
 
 recordConstructor  'record'
   = '{' _nl args:array_recordArguments _nl '}' 
-  { return constructRecord(args); }
+  { return h.constructRecord(args); }
 
   array_recordArguments 
     = head:keyValuePair tail:(_ ',' _nl @keyValuePair)* 
@@ -338,7 +271,7 @@ recordConstructor  'record'
 
   keyValuePair 
     = key:expression _ ':' _nl value:expression 
-    { return nodeKeyValue(key, value)}
+    { return h.nodeKeyValue(key, value)}
 
 // Separators
 
@@ -377,30 +310,30 @@ statementSeparator 'statement separator'
 noArguments = ('(' _nl ')' )?
 
 typeIdentifier 'type identifier'
-  = ([a-z]+[_a-z0-9]i*) {return nodeTypeIdentifier(text())} 
+  = ([a-z]+[_a-z0-9]i*) {return h.nodeTypeIdentifier(text())} 
 
 typeConstructorIdentifier 'type constructor identifier'
-  = ([A-Z]+[_a-z0-9]i*) {return nodeTypeIdentifier(text())} 
+  = ([A-Z]+[_a-z0-9]i*) {return h.nodeTypeIdentifier(text())} 
 
 typeExpression = typePostModifierExpression
 
 typePostModifierExpression = head:typeOr tail:(_ '$' _nl @typeModifier)*
   { 
     return tail.reduce((result, element) => {
-      return makeFunctionCall('$_typeModifier_'+element.modifier.value+'_$', [result, ...element.args])
+      return h.makeFunctionCall('$_typeModifier_'+element.modifier.value+'_$', [result, ...element.args])
     }, head)
   }
 
 typeOr = head:typeFunction tail:(_ '|' _nl @typeFunction)*
-  { return tail.length === 0 ? head : apply('$_typeOr_$', constructArray([head, ...tail])); }
+  { return tail.length === 0 ? head : h.apply('$_typeOr_$', h.constructArray([head, ...tail])); }
 
 typeFunction = head:typeModifierExpression tail:(_ '=>' _nl @typeModifierExpression)*    
-  { return tail.length === 0 ? head : apply( '$_typeFunction_$', constructArray([head, ...tail])); }
+  { return tail.length === 0 ? head : h.apply( '$_typeFunction_$', h.constructArray([head, ...tail])); }
 
 typeModifierExpression = head:basicType tail:(_ '<-' _nl @typeModifier)* 
   { 
     return tail.reduce((result, element) => {
-      return makeFunctionCall('$_typeModifier_'+element.modifier.value+'_$', [result, ...element.args])
+      return h.makeFunctionCall('$_typeModifier_'+element.modifier.value+'_$', [result, ...element.args])
     }, head)
   }
 
@@ -413,10 +346,10 @@ typeModifierExpression = head:basicType tail:(_ '<-' _nl @typeModifier)*
 basicType = typeConstructor / typeArray / typeRecord / typeInParanthesis / typeIdentifier
 
 typeArray = '[' _nl elem:typeExpression _nl ']' 
-  {return apply('$_typeArray_$', elem)}
+  {return h.apply('$_typeArray_$', elem)}
 
 typeRecord = '{' _nl elems:array_typeRecordArguments _nl '}' 
-  { return apply('$_typeRecord_$', constructRecord(elems)); }  
+  { return h.apply('$_typeRecord_$', h.constructRecord(elems)); }  
 
   array_typeRecordArguments 
     = head:typeKeyValuePair tail:(_ ',' _nl @typeKeyValuePair)* 
@@ -424,22 +357,22 @@ typeRecord = '{' _nl elems:array_typeRecordArguments _nl '}'
 
   typeKeyValuePair 
     = key:identifier _ ':' _nl value:typeExpression 
-    { return nodeKeyValue(key, value)}
+    { return h.nodeKeyValue(key, value)}
 
 typeConstructor 
   = constructor:typeConstructorIdentifier _ '(' _nl args:array_types _nl ')' 
-    { return makeFunctionCall('$_typeConstructor_$', [constructor, constructArray(args)]); }
+    { return h.makeFunctionCall('$_typeConstructor_$', [constructor, h.constructArray(args)]); }
   / constructor:typeConstructorIdentifier _ noArguments
-    { return makeFunctionCall('$_typeConstructor_$', [constructor, constructArray([])]); }
+    { return h.makeFunctionCall('$_typeConstructor_$', [constructor, h.constructArray([])]); }
 
   array_types = head:typeExpression tail:(_ ',' _nl @typeExpression)* 
     { return [head, ...tail]; }
 
 typeStatement = typeAliasStatement / typeOfStatement
 typeAliasStatement = 'type' __nl typeIdentifier:typeIdentifier _nl '=' _nl typeExpression:typeExpression 
-  { return makeFunctionCall('$_typeAlias_$', [typeIdentifier, typeExpression])}
+  { return h.makeFunctionCall('$_typeAlias_$', [typeIdentifier, typeExpression])}
 typeOfStatement = identifier:identifier _ ':' _nl typeExpression:typeExpression
-  { return makeFunctionCall('$_typeOf_$', [identifier, typeExpression])}
+  { return h.makeFunctionCall('$_typeOf_$', [identifier, typeExpression])}
 
 typeInParanthesis = '(' _nl typeExpression:typeExpression _nl ')' {return typeExpression}
 
@@ -447,4 +380,4 @@ typeInParanthesis = '(' _nl typeExpression:typeExpression _nl ')' {return typeEx
 // TODO: Example of foo = {a: 2, b: 5}; type fooKeys = string $ memberOf(foo->keys)
 // TODO: Example of memberOf( [1,2,3] )
 // TODO: Example of $ 
-// TODO: Cons(a, list) | EmptyList
+// TODO: Cons(a, list) | EmptyList

packages/squiggle-lang/src/rescript/Reducer/Reducer_Peggy/helpers.ts

@@ -0,0 +1,215 @@
+import { LocationRange } from "peggy";
+
+export const toFunction = {
+  "-": "subtract",
+  "->": "pipe",
+  "!=": "unequal",
+  ".-": "dotSubtract",
+  ".*": "dotMultiply",
+  "./": "dotDivide",
+  ".^": "dotPow",
+  ".+": "dotAdd",
+  "*": "multiply",
+  "/": "divide",
+  "&&": "and",
+  "^": "pow", // or xor
+  "+": "add",
+  "<": "smaller",
+  "<=": "smallerEq",
+  "==": "equal",
+  ">": "larger",
+  ">=": "largerEq",
+  "||": "or",
+  to: "credibleIntervalToDistribution",
+};
+
+export const unaryToFunction = {
+  "-": "unaryMinus",
+  "!": "not",
+  ".-": "unaryDotMinus",
+};
+
+export const postOperatorToFunction = {
+  ".": "$_atIndex_$",
+  "()": "$$_applyAll_$$",
+  "[]": "$_atIndex_$",
+};
+
+type NodeBlock = {
+  type: "Block";
+  statements: AnyPeggyNode[];
+};
+
+type NodeExpression = {
+  type: "Expression";
+  nodes: AnyPeggyNode[];
+};
+
+type NodeFloat = {
+  type: "Float";
+  value: number;
+};
+
+type NodeInteger = {
+  type: "Integer";
+  value: number;
+};
+
+type NodeIdentifier = {
+  type: "Identifier";
+  value: string;
+  location: LocationRange;
+};
+
+type NodeCallIdentifier = {
+  type: "CallIdentifier";
+  value: string;
+};
+
+type NodeLetStatement = {
+  type: "LetStatement";
+  variable: NodeIdentifier;
+  value: AnyPeggyNode;
+};
+
+type NodeLambda = {
+  type: "Lambda";
+  args: AnyPeggyNode[];
+  body: AnyPeggyNode;
+};
+
+type NodeTernary = {
+  type: "Ternary";
+  condition: AnyPeggyNode;
+  trueExpression: AnyPeggyNode;
+  falseExpression: AnyPeggyNode;
+};
+
+type NodeKeyValue = {
+  type: "KeyValue";
+  key: AnyPeggyNode;
+  value: AnyPeggyNode;
+};
+
+type NodeString = {
+  type: "String";
+  value: string;
+  location?: LocationRange;
+};
+
+type NodeBoolean = {
+  type: "Boolean";
+  value: boolean;
+};
+
+export type AnyPeggyNode =
+  | NodeBlock
+  | NodeExpression
+  | NodeFloat
+  | NodeInteger
+  | NodeIdentifier
+  | NodeCallIdentifier
+  | NodeLetStatement
+  | NodeLambda
+  | NodeTernary
+  | NodeKeyValue
+  | NodeString
+  | NodeBoolean;
+
+export function makeFunctionCall(fn: string, args: AnyPeggyNode[]) {
+  if (fn === "$$_applyAll_$$") {
+    // Any list of values is applied from left to right anyway.
+    // Like in Haskell and Lisp.
+    // So we remove the redundant $$_applyAll_$$.
+    if (args[0].type === "Identifier") {
+      args[0] = {
+        ...args[0],
+        type: "CallIdentifier",
+      };
+    }
+    return nodeExpression(args);
+  } else {
+    return nodeExpression([nodeCallIndentifier(fn), ...args]);
+  }
+}
+
+export function apply(fn: string, arg: AnyPeggyNode) {
+  return makeFunctionCall(fn, [arg]);
+}
+export function constructArray(elems: AnyPeggyNode[]) {
+  return apply("$_constructArray_$", nodeExpression(elems));
+}
+export function constructRecord(elems: AnyPeggyNode[]) {
+  return apply("$_constructRecord_$", nodeExpression(elems));
+}
+
+export function nodeBlock(statements: AnyPeggyNode[]): NodeBlock {
+  return { type: "Block", statements };
+}
+export function nodeBoolean(value: boolean): NodeBoolean {
+  return { type: "Boolean", value };
+}
+export function nodeCallIndentifier(value: string): NodeCallIdentifier {
+  return { type: "CallIdentifier", value };
+}
+export function nodeExpression(args: AnyPeggyNode[]): NodeExpression {
+  return { type: "Expression", nodes: args };
+}
+export function nodeFloat(value: number): NodeFloat {
+  return { type: "Float", value };
+}
+export function nodeIdentifier(
+  value: string,
+  location: LocationRange
+): NodeIdentifier {
+  return { type: "Identifier", value, location };
+}
+export function nodeInteger(value: number): NodeInteger {
+  return { type: "Integer", value };
+}
+export function nodeKeyValue(
+  key: AnyPeggyNode,
+  value: AnyPeggyNode
+): NodeKeyValue {
+  if (key.type === "Identifier") {
+    key = {
+      ...key,
+      type: "String",
+    };
+  }
+  return { type: "KeyValue", key, value };
+}
+export function nodeLambda(
+  args: AnyPeggyNode[],
+  body: AnyPeggyNode
+): NodeLambda {
+  return { type: "Lambda", args, body };
+}
+export function nodeLetStatement(
+  variable: NodeIdentifier,
+  value: AnyPeggyNode
+): NodeLetStatement {
+  return { type: "LetStatement", variable, value };
+}
+export function nodeModuleIdentifier(value: string) {
+  return { type: "ModuleIdentifier", value };
+}
+export function nodeString(value: string): NodeString {
+  return { type: "String", value };
+}
+export function nodeTernary(
+  condition: AnyPeggyNode,
+  trueExpression: AnyPeggyNode,
+  falseExpression: AnyPeggyNode
+): NodeTernary {
+  return {
+    type: "Ternary",
+    condition,
+    trueExpression,
+    falseExpression,
+  };
+}
+
+export function nodeTypeIdentifier(typeValue: string) {
+  return { type: "TypeIdentifier", value: typeValue };
+}

packages/vscode-ext/client/src/extension.ts

@@ -4,6 +4,7 @@ import * as vscode from "vscode";
 import { startClient, stopClient } from "./client";
 
 import { SquiggleEditorProvider } from "./editor";
+import { registerSemanticHighlight } from "./highlight";
 import { registerPreviewCommand } from "./preview";
 
 // this method is called when your extension is activated
@@ -13,6 +14,8 @@ export function activate(context: vscode.ExtensionContext) {
 
   registerPreviewCommand(context);
 
+  registerSemanticHighlight();
+
   startClient(context);
 }
 

packages/vscode-ext/client/src/highlight.ts

@@ -0,0 +1,96 @@
+import * as vscode from "vscode";
+
+import { parse } from "@quri/squiggle-lang";
+import { AnyPeggyNode } from "@quri/squiggle-lang/dist/src/rescript/Reducer/Reducer_Peggy/helpers";
+
+const tokenTypes = ["enum", "function", "variable", "property"];
+const tokenModifiers = ["declaration", "documentation"];
+const legend = new vscode.SemanticTokensLegend(tokenTypes, tokenModifiers);
+
+const convertRange = (
+  range: Extract<AnyPeggyNode, { type: "Identifier" }>["location"]
+) =>
+  new vscode.Range(
+    new vscode.Position(range.start.line - 1, range.start.column - 1),
+    new vscode.Position(range.end.line - 1, range.end.column - 1)
+  );
+
+const populateTokensBuilder = (
+  tokensBuilder: vscode.SemanticTokensBuilder,
+  node: AnyPeggyNode
+  // bindings: { [key: string]: boolean }
+) => {
+  switch (node.type) {
+    case "Expression":
+      for (const child of node.nodes) {
+        populateTokensBuilder(tokensBuilder, child);
+      }
+      break;
+    case "Block":
+      for (const child of node.statements) {
+        populateTokensBuilder(tokensBuilder, child);
+      }
+      break;
+    case "LetStatement":
+      tokensBuilder.push(
+        convertRange(node.variable.location),
+        node.value.type === "Lambda" ? "function" : "variable",
+        ["declaration"]
+      );
+      populateTokensBuilder(tokensBuilder, node.value);
+      break;
+    case "Lambda":
+      for (const arg of node.args) {
+        populateTokensBuilder(tokensBuilder, arg);
+      }
+      populateTokensBuilder(tokensBuilder, node.body);
+      break;
+    case "Ternary":
+      populateTokensBuilder(tokensBuilder, node.condition);
+      populateTokensBuilder(tokensBuilder, node.trueExpression);
+      populateTokensBuilder(tokensBuilder, node.falseExpression);
+      break;
+    case "KeyValue":
+      if (node.key.type === "String" && node.key.location) {
+        tokensBuilder.push(convertRange(node.key.location), "property", [
+          "declaration",
+        ]);
+      } else {
+        populateTokensBuilder(tokensBuilder, node.key);
+      }
+      populateTokensBuilder(tokensBuilder, node.value);
+      break;
+    case "Identifier":
+      tokensBuilder.push(convertRange(node.location), "variable");
+      break;
+  }
+};
+
+export const registerSemanticHighlight = () => {
+  const provider: vscode.DocumentSemanticTokensProvider = {
+    provideDocumentSemanticTokens(
+      document: vscode.TextDocument
+    ): vscode.ProviderResult<vscode.SemanticTokens> {
+      const parseResult = parse(document.getText());
+
+      const tokensBuilder = new vscode.SemanticTokensBuilder(legend);
+      if (parseResult.tag === "Ok") {
+        populateTokensBuilder(
+          tokensBuilder,
+          parseResult.value
+          // {}
+        );
+      }
+
+      return tokensBuilder.build();
+    },
+  };
+
+  const selector = { language: "squiggle", scheme: "file" };
+
+  vscode.languages.registerDocumentSemanticTokensProvider(
+    selector,
+    provider,
+    legend
+  );
+};

packages/vscode-ext/server/src/server.ts

@@ -13,6 +13,10 @@ import { parse } from "@quri/squiggle-lang";
 
 import { TextDocument } from "vscode-languageserver-textdocument";
 
+// Documentation:
+// - https://code.visualstudio.com/api/language-extensions/language-server-extension-guide
+// - https://microsoft.github.io/language-server-protocol/specifications/specification-current
+
 // Create a connection for the server, using Node's IPC as a transport.
 // Also include all preview / proposed LSP features.
 let connection = createConnection(ProposedFeatures.all);
@@ -23,17 +27,7 @@ documents.onDidChangeContent((change) => {
   validateSquiggleDocument(change.document);
 });
 
-let hasDiagnosticRelatedInformationCapability = false;
-
 connection.onInitialize((params: InitializeParams) => {
-  const capabilities = params.capabilities;
-
-  hasDiagnosticRelatedInformationCapability = !!(
-    capabilities.textDocument &&
-    capabilities.textDocument.publishDiagnostics &&
-    capabilities.textDocument.publishDiagnostics.relatedInformation
-  );
-
   const result: InitializeResult = {
     capabilities: {
       textDocumentSync: TextDocumentSyncKind.Incremental,