Reimplement 'babel-template' with better caching and tagged literal utilities.

packages/babel-template/src/builder.js

@@ -0,0 +1,131 @@
+// @flow
+
+import {
+  merge,
+  validate,
+  type TemplateOpts,
+  type PublicOpts,
+  type PublicReplacements,
+} from "./options";
+import type { Formatter } from "./formatters";
+
+import stringTemplate from "./string";
+import literalTemplate from "./literal";
+
+export type TemplateBuilder<T> = {
+  // Build a new builder, merging the given options with the previous ones.
+  (opts: PublicOpts): TemplateBuilder<T>,
+
+  // Building from a string produces an AST builder function by default.
+  (tpl: string, opts: ?PublicOpts): (?PublicReplacements) => T,
+
+  // Building from a template literal produces an AST builder function by default.
+  (tpl: Array<string>, ...args: Array<mixed>): (?PublicReplacements) => T,
+
+  // Allow users to explicitly create templates that produce ASTs, skipping
+  // the need for an intermediate function.
+  ast: {
+    (tpl: string, opts: ?PublicOpts): T,
+    (tpl: Array<string>, ...args: Array<mixed>): T,
+  },
+};
+
+// Prebuild the options that will be used when parsing a `.ast` template.
+// These do not use a pattern because there is no way for users to pass in
+// replacement patterns to begin with, and disabling pattern matching means
+// users have more flexibility in what type of content they have in their
+// template JS.
+const NO_PLACEHOLDER: TemplateOpts = validate({
+  placeholderPattern: false,
+});
+
+export default function createTemplateBuilder<T>(
+  formatter: Formatter<T>,
+  defaultOpts?: TemplateOpts,
+): TemplateBuilder<T> {
+  const templateFnCache = new WeakMap();
+  const templateAstCache = new WeakMap();
+  const cachedOpts = defaultOpts || validate(null);
+
+  return Object.assign(
+    ((tpl, ...args) => {
+      if (typeof tpl === "string") {
+        if (args.length > 1) throw new Error("Unexpected extra params.");
+        return extendedTrace(
+          stringTemplate(formatter, tpl, merge(cachedOpts, validate(args[0]))),
+        );
+      } else if (Array.isArray(tpl)) {
+        let builder = templateFnCache.get(tpl);
+        if (!builder) {
+          builder = literalTemplate(formatter, tpl, cachedOpts);
+          templateFnCache.set(tpl, builder);
+        }
+        return extendedTrace(builder(args));
+      } else if (typeof tpl === "object" && tpl) {
+        if (args.length > 0) throw new Error("Unexpected extra params.");
+        return createTemplateBuilder(
+          formatter,
+          merge(cachedOpts, validate(tpl)),
+        );
+      }
+      throw new Error(`Unexpected template param ${typeof tpl}`);
+    }: Function),
+    {
+      ast: (tpl, ...args) => {
+        if (typeof tpl === "string") {
+          if (args.length > 1) throw new Error("Unexpected extra params.");
+          return stringTemplate(
+            formatter,
+            tpl,
+            merge(merge(cachedOpts, validate(args[0])), NO_PLACEHOLDER),
+          )();
+        } else if (Array.isArray(tpl)) {
+          let builder = templateAstCache.get(tpl);
+          if (!builder) {
+            builder = literalTemplate(
+              formatter,
+              tpl,
+              merge(cachedOpts, NO_PLACEHOLDER),
+            );
+            templateAstCache.set(tpl, builder);
+          }
+          return builder(args)();
+        }
+
+        throw new Error(`Unexpected template param ${typeof tpl}`);
+      },
+    },
+  );
+}
+
+function extendedTrace<Arg, Result>(fn: Arg => Result): Arg => Result {
+  // Since we lazy parse the template, we get the current stack so we have the
+  // original stack to append if it errors when parsing
+  let rootStack = "";
+  try {
+    // error stack gets populated in IE only on throw
+    // (https://msdn.microsoft.com/en-us/library/hh699850(v=vs.94).aspx)
+    throw new Error();
+  } catch (error) {
+    if (error.stack) {
+      // error.stack does not exists in IE <= 9
+      // We slice off the top 3 items in the stack to remove the call to
+      // 'extendedTrace', and the anonymous builder function, with the final
+      // stripped line being the error message itself since we threw it
+      // in the first place and it doesn't matter.
+      rootStack = error.stack
+        .split("\n")
+        .slice(3)
+        .join("\n");
+    }
+  }
+
+  return (arg: Arg) => {
+    try {
+      return fn(arg);
+    } catch (err) {
+      err.stack += `\n    =============\n${rootStack}`;
+      throw err;
+    }
+  };
+}

packages/babel-template/src/formatters.js

@@ -0,0 +1,74 @@
+// @flow
+
+export type Formatter<T> = {
+  code: string => string,
+  validate: BabelNodeFile => void,
+  unwrap: BabelNodeFile => T,
+};
+
+function makeStatementFormatter<T>(
+  fn: (Array<BabelNodeStatement>) => T,
+): Formatter<T> {
+  return {
+    // We need to prepend a ";" to force statement parsing so that
+    // ExpressionStatement strings won't be parsed as directives.
+    // Alonside that, we also prepend a comment so that when a syntax error
+    // is encountered, the user will be less likely to get confused about
+    // where the random semicolon came from.
+    code: str => `/* @babel/template */;\n${str}`,
+    validate: () => {},
+    unwrap: (ast: BabelNodeFile): T => {
+      return fn(ast.program.body.slice(1));
+    },
+  };
+}
+
+export const smart: Formatter<
+  Array<BabelNodeStatement> | BabelNodeStatement,
+> = makeStatementFormatter(body => {
+  if (body.length > 1) {
+    return body;
+  } else {
+    return body[0];
+  }
+});
+
+export const statements: Formatter<
+  Array<BabelNodeStatement>,
+> = makeStatementFormatter(body => body);
+
+export const statement: Formatter<
+  BabelNodeStatement,
+> = makeStatementFormatter(body => {
+  // We do this validation when unwrapping since the replacement process
+  // could have added or removed statements.
+  if (body.length === 0) {
+    throw new Error("Found nothing to return.");
+  }
+  if (body.length > 1) {
+    throw new Error("Found multiple statements but wanted one");
+  }
+
+  return body[0];
+});
+
+export const expression: Formatter<BabelNodeExpression> = {
+  code: str => `(\n${str}\n)`,
+  validate: (ast: BabelNodeFile) => {
+    const { program } = ast;
+    if (program.body.length > 1) {
+      throw new Error("Found multiple statements but wanted one");
+    }
+    const expression = program.body[0].expression;
+    if (expression.start === 0) {
+      throw new Error("Parse result included parens.");
+    }
+  },
+  unwrap: ast => ast.program.body[0].expression,
+};
+
+export const program: Formatter<BabelNodeProgram> = {
+  code: str => str,
+  validate: () => {},
+  unwrap: ast => ast.program,
+};

packages/babel-template/src/index.js

@@ -1,10 +1,31 @@
-import factory from "./string";
-import template from "./literal";
+// @flow
 
-export default function(firstArg, ...rest) {
-  if (typeof firstArg === "string") {
-    return factory(firstArg, ...rest);
-  } else {
-    return template(firstArg, ...rest);
-  }
-}
+import * as formatters from "./formatters";
+import createTemplateBuilder from "./builder";
+
+export const smart = createTemplateBuilder(formatters.smart);
+export const statement = createTemplateBuilder(formatters.statement);
+export const statements = createTemplateBuilder(formatters.statements);
+export const expression = createTemplateBuilder(formatters.expression);
+export const program = createTemplateBuilder(formatters.program);
+
+type DefaultTemplateBuilder = typeof smart & {
+  smart: typeof smart,
+  statement: typeof statement,
+  statements: typeof statements,
+  expression: typeof expression,
+  program: typeof program,
+  ast: typeof smart.ast,
+};
+
+export default Object.assign(
+  ((smart.bind(undefined): any): DefaultTemplateBuilder),
+  {
+    smart,
+    statement,
+    statements,
+    expression,
+    program,
+    ast: smart.ast,
+  },
+);

packages/babel-template/src/literal.js

@@ -1,97 +1,103 @@
-import * as t from "@babel/types";
+// @flow
 
-import factory from "./string";
+import type { Formatter } from "./formatters";
+import { normalizeReplacements, type TemplateOpts } from "./options";
+import parseAndBuildMetadata from "./parse";
+import populatePlaceholders from "./populate";
 
-export default function template(
-  partials: Object | string[],
-  ...args: Array<Object>
-) {
-  if (!Array.isArray(partials)) {
-    // support template({ options })`string`
-    return templateApply.bind(undefined, partials);
-  }
-  return templateApply(null, partials, ...args);
-}
+export default function literalTemplate<T>(
+  formatter: Formatter<T>,
+  tpl: Array<string>,
+  opts: TemplateOpts,
+): (Array<mixed>) => mixed => T {
+  const { metadata, names } = buildLiteralData(formatter, tpl, opts);
 
-function templateApply(
-  opts: Object | null,
-  partials: string[],
-  ...args: Array<Object>
-) {
-  if (partials.some(str => str.includes("$BABEL_TEMPLATE$"))) {
-    throw new Error("Template contains illegal substring $BABEL_TEMPLATE$");
-  }
+  return (arg: Array<mixed>) => {
+    const defaultReplacements = arg.reduce((acc, replacement, i) => {
+      acc[names[i]] = replacement;
+      return acc;
+    }, {});
 
-  if (partials.length == 1) {
-    return factory(partials[0], opts);
-  }
+    return (arg: mixed) => {
+      const replacements = normalizeReplacements(arg);
 
-  const replacementSet = new Set();
-  const replacementMap = new Map();
-  const replacementValueMap = new Map();
-  let hasNonNumericReplacement = false;
-  for (const arg of args) {
-    if (replacementMap.has(arg)) {
-      continue;
-    }
-
-    if (typeof arg === "number") {
-      replacementMap.set(arg, `$${arg}`);
-    } else if (typeof arg === "string") {
-      // avoid duplicates should t.toIdentifier produce the same result for different arguments
-      const replacementBase = `$BABEL_TEMPLATE$$${t.toIdentifier(arg)}`;
-      let replacement = replacementBase;
-      for (let i = 2; replacementSet.has(replacement); i++) {
-        replacement = `${replacementBase}${i}`;
-      }
-      replacementSet.add(replacement);
-      replacementMap.set(arg, replacement);
-      hasNonNumericReplacement = true;
-    } else {
-      // there can't be duplicates as the size always grows
-      const name = `$BABEL_TEMPLATE$VALUE$${replacementValueMap.size}`;
-
-      // TODO: check if the arg is a Node
-      replacementMap.set(arg, name);
-      replacementValueMap.set(name, arg);
-      hasNonNumericReplacement = true;
-    }
-  }
-
-  if (hasNonNumericReplacement && replacementMap.has(0)) {
-    throw new Error(
-      "Template cannot have a '0' replacement and a named replacement at the same time",
-    );
-  }
-
-  const code = partials.reduce((acc, partial, i) => {
-    if (acc == null) {
-      return partial;
-    }
-
-    const replacement = replacementMap.get(args[i - 1]);
-    return `${acc}${replacement}${partial}`;
-  }, null);
-
-  const func = factory(code, opts);
-
-  return (...args: Array<Object>) => {
-    if (hasNonNumericReplacement) {
-      const argObj = args[0] || {};
-      const converted = {};
-
-      for (const [key, replacement] of replacementMap) {
-        if (typeof key === "number") continue;
-        if (replacementValueMap.has(replacement)) {
-          converted[replacement] = replacementValueMap.get(replacement);
-        } else {
-          converted[replacement] = argObj[key];
-        }
+      if (replacements) {
+        Object.keys(replacements).forEach(key => {
+          if (Object.prototype.hasOwnProperty.call(defaultReplacements, key)) {
+            throw new Error("Unexpected replacement overlap.");
+          }
+        });
       }
 
-      args[0] = converted;
-    }
-
-    return func(...args);
+      return formatter.unwrap(
+        populatePlaceholders(
+          metadata,
+          replacements
+            ? Object.assign(replacements, defaultReplacements)
+            : defaultReplacements,
+        ),
+      );
+    };
   };
 }
+
+function buildLiteralData<T>(
+  formatter: Formatter<T>,
+  tpl: Array<string>,
+  opts: TemplateOpts,
+) {
+  let names;
+  let nameSet;
+  let metadata;
+  let prefix = "";
+
+  do {
+    // If there are cases where the template already contains $0 or any other
+    // matching pattern, we keep adding "$" characters until a unique prefix
+    // is found.
+    prefix += "$";
+    const result = buildTemplateCode(tpl, prefix);
+
+    names = result.names;
+    nameSet = new Set(names);
+    metadata = parseAndBuildMetadata(formatter, formatter.code(result.code), {
+      parser: opts.parser,
+
+      // Explicitly include our generated names in the whitelist so users never
+      // have to think about whether their placeholder pattern will match.
+      placeholderWhitelist: new Set(
+        result.names.concat(
+          opts.placeholderWhitelist
+            ? Array.from(opts.placeholderWhitelist)
+            : [],
+        ),
+      ),
+      placeholderPattern: opts.placeholderPattern,
+      preserveComments: opts.preserveComments,
+    });
+  } while (
+    metadata.placeholders.some(
+      placeholder => placeholder.isDuplicate && nameSet.has(placeholder.name),
+    )
+  );
+
+  return { metadata, names };
+}
+
+function buildTemplateCode(
+  tpl: Array<string>,
+  prefix: string,
+): { names: Array<string>, code: string } {
+  const names = [];
+
+  let code = tpl[0];
+
+  for (let i = 1; i < tpl.length; i++) {
+    const value = `${prefix}${i - 1}`;
+    names.push(value);
+
+    code += value + tpl[i];
+  }
+
+  return { names, code };
+}

packages/babel-template/src/options.js

@@ -0,0 +1,116 @@
+// @flow
+
+/**
+ * These are the options that 'babel-template' actually accepts and typechecks
+ * when called. All other options are passed through to the parser.
+ */
+export type PublicOpts = {
+  /**
+   * A set of placeholder names to automatically accept, ignoring the given
+   * pattern entirely.
+   */
+  placeholderWhitelist?: ?Set<string>,
+
+  /**
+   * A pattern to search for when looking for Identifier and StringLiteral
+   * nodes that can be replaced.
+   *
+   * 'false' will disable placeholder searching entirely, leaving only the
+   * 'placeholderWhitelist' value to find replacements.
+   *
+   * Defaults to /^[_$A-Z0-9]+$/.
+   */
+  placeholderPattern?: ?(RegExp | false),
+
+  /**
+   * 'true' to pass through comments from the template into the resulting AST,
+   * or 'false' to automatically discard comments. Defaults to 'false'.
+   */
+  preserveComments?: ?boolean,
+};
+
+export type TemplateOpts = {|
+  parser: {},
+  placeholderWhitelist: Set<string> | void,
+  placeholderPattern: RegExp | false | void,
+  preserveComments: boolean | void,
+|};
+
+export function merge(a: TemplateOpts, b: TemplateOpts): TemplateOpts {
+  const {
+    placeholderWhitelist = a.placeholderWhitelist,
+    placeholderPattern = a.placeholderPattern,
+    preserveComments = a.preserveComments,
+  } = b;
+
+  return {
+    parser: Object.assign({}, a.parser, b.parser),
+    placeholderWhitelist,
+    placeholderPattern,
+    preserveComments,
+  };
+}
+
+export function validate(opts: mixed): TemplateOpts {
+  if (opts != null && typeof opts !== "object") {
+    throw new Error("Unknown template options.");
+  }
+
+  const {
+    placeholderWhitelist,
+    placeholderPattern,
+    preserveComments,
+    ...parser
+  } =
+    opts || {};
+
+  if (placeholderWhitelist != null && !(placeholderWhitelist instanceof Set)) {
+    throw new Error(
+      "'.placeholderWhitelist' must be a Set, null, or undefined",
+    );
+  }
+
+  if (
+    placeholderPattern != null &&
+    !(placeholderPattern instanceof RegExp) &&
+    placeholderPattern !== false
+  ) {
+    throw new Error(
+      "'.placeholderPattern' must be a RegExp, false, null, or undefined",
+    );
+  }
+
+  if (preserveComments != null && typeof preserveComments !== "boolean") {
+    throw new Error(
+      "'.preserveComments' must be a boolean, null, or undefined",
+    );
+  }
+
+  return {
+    parser,
+    placeholderWhitelist: placeholderWhitelist || undefined,
+    placeholderPattern:
+      placeholderPattern == null ? undefined : placeholderPattern,
+    preserveComments: preserveComments == null ? false : preserveComments,
+  };
+}
+
+export type PublicReplacements = { [string]: mixed } | Array<mixed>;
+export type TemplateReplacements = { [string]: mixed } | void;
+
+export function normalizeReplacements(
+  replacements: mixed,
+): TemplateReplacements {
+  if (Array.isArray(replacements)) {
+    return replacements.reduce((acc, replacement, i) => {
+      acc["$" + i] = replacement;
+      return acc;
+    }, {});
+  } else if (typeof replacements === "object" || replacements == null) {
+    return replacements || undefined;
+  }
+
+  throw new Error(
+    "Template replacements must be an array, object, null, or undefined",
+  );
+}

packages/babel-template/src/parse.js

@@ -0,0 +1,156 @@
+// @flow
+import * as t from "@babel/types";
+import type { TraversalAncestors, TraversalHandler } from "@babel/types";
+import { parse } from "babylon";
+import { codeFrameColumns } from "@babel/code-frame";
+import type { TemplateOpts } from "./options";
+import type { Formatter } from "./formatters";
+
+export type Metadata = {
+  ast: BabelNodeFile,
+  placeholders: Array<Placeholder>,
+  placeholderNames: Set<string>,
+};
+
+type PlaceholderType = "string" | "param" | "statement" | "other";
+export type Placeholder = {|
+  name: string,
+  resolve: BabelNodeFile => { parent: BabelNode, key: string, index?: number },
+  type: PlaceholderType,
+  isDuplicate: boolean,
+|};
+
+const PATTERN = /^[_$A-Z0-9]+$/;
+
+export default function parseAndBuildMetadata<T>(
+  formatter: Formatter<T>,
+  code: string,
+  opts: TemplateOpts,
+): Metadata {
+  const ast = parseWithCodeFrame(code, opts.parser);
+
+  const {
+    placeholderWhitelist,
+    placeholderPattern = PATTERN,
+    preserveComments,
+  } = opts;
+
+  t.removePropertiesDeep(ast, {
+    preserveComments,
+  });
+
+  formatter.validate(ast);
+
+  const placeholders = [];
+  const placeholderNames = new Set();
+
+  t.traverse(ast, (placeholderVisitorHandler: TraversalHandler<*>), {
+    placeholders,
+    placeholderNames,
+    placeholderWhitelist,
+    placeholderPattern,
+  });
+
+  return {
+    ast,
+    placeholders,
+    placeholderNames,
+  };
+}
+
+function placeholderVisitorHandler(
+  node: BabelNode,
+  ancestors: TraversalAncestors,
+  state: MetadataState,
+) {
+  let name;
+  if (t.isIdentifier(node)) {
+    name = ((node: any): BabelNodeIdentifier).name;
+  } else if (t.isStringLiteral(node)) {
+    name = ((node: any): BabelNodeStringLiteral).value;
+  } else {
+    return;
+  }
+
+  if (
+    (!state.placeholderPattern || !state.placeholderPattern.test(name)) &&
+    (!state.placeholderWhitelist || !state.placeholderWhitelist.has(name))
+  ) {
+    return;
+  }
+
+  // Keep our own copy of the ancestors so we can use it in .resolve().
+  ancestors = ancestors.slice();
+
+  const { node: parent, key } = ancestors[ancestors.length - 1];
+
+  let type: PlaceholderType;
+  if (t.isStringLiteral(node)) {
+    type = "string";
+  } else if (
+    (t.isNewExpression(parent) && key === "arguments") ||
+    (t.isCallExpression(parent) && key === "arguments") ||
+    (t.isFunction(parent) && key === "params")
+  ) {
+    type = "param";
+  } else if (t.isExpressionStatement(parent)) {
+    type = "statement";
+    ancestors = ancestors.slice(0, -1);
+  } else {
+    type = "other";
+  }
+
+  state.placeholders.push({
+    name,
+    type,
+    resolve: ast => resolveAncestors(ast, ancestors),
+    isDuplicate: state.placeholderNames.has(name),
+  });
+  state.placeholderNames.add(name);
+}
+
+function resolveAncestors(ast: BabelNodeFile, ancestors: TraversalAncestors) {
+  let parent: BabelNode = ast;
+  for (let i = 0; i < ancestors.length - 1; i++) {
+    const { key, index } = ancestors[i];
+
+    if (index === undefined) {
+      parent = (parent: any)[key];
+    } else {
+      parent = (parent: any)[key][index];
+    }
+  }
+
+  const { key, index } = ancestors[ancestors.length - 1];
+
+  return { parent, key, index };
+}
+
+type MetadataState = {
+  placeholders: Array<Placeholder>,
+  placeholderNames: Set<string>,
+  placeholderWhitelist: Set<string> | void,
+  placeholderPattern: RegExp | false,
+};
+
+function parseWithCodeFrame(code: string, parserOpts: {}): BabelNodeFile {
+  parserOpts = Object.assign(
+    {
+      allowReturnOutsideFunction: true,
+      allowSuperOutsideMethod: true,
+      sourceType: "module",
+    },
+    parserOpts,
+  );
+
+  try {
+    return parse(code, parserOpts);
+  } catch (err) {
+    const loc = err.loc;
+    if (loc) {
+      err.loc = null;
+      err.message += "\n" + codeFrameColumns(code, { start: loc });
+    }
+    throw err;
+  }
+}

packages/babel-template/src/populate.js

@@ -0,0 +1,131 @@
+// @flow
+import * as t from "@babel/types";
+
+import type { TemplateReplacements } from "./options";
+import type { Metadata, Placeholder } from "./parse";
+
+export default function populatePlaceholders(
+  metadata: Metadata,
+  replacements: TemplateReplacements,
+): BabelNodeFile {
+  const ast = t.cloneDeep(metadata.ast);
+
+  if (replacements) {
+    metadata.placeholders.forEach(placeholder => {
+      if (
+        !Object.prototype.hasOwnProperty.call(replacements, placeholder.name)
+      ) {
+        throw new Error(`No substitution given for "${placeholder.name}"`);
+      }
+    });
+    Object.keys(replacements).forEach(key => {
+      if (!metadata.placeholderNames.has(key)) {
+        throw new Error(`Unknown substitution "${key}" given`);
+      }
+    });
+  }
+
+  // Process in reverse order to AST mutation doesn't change indices that
+  // will be needed for later calls to `placeholder.resolve()`.
+  metadata.placeholders
+    .slice()
+    .reverse()
+    .forEach(placeholder => {
+      try {
+        applyReplacement(
+          placeholder,
+          ast,
+          (replacements && replacements[placeholder.name]) || null,
+        );
+      } catch (e) {
+        e.message = `babel-template placeholder "${placeholder.name}": ${e.message}`;
+        throw e;
+      }
+    });
+
+  return ast;
+}
+
+function applyReplacement(
+  placeholder: Placeholder,
+  ast: BabelNodeFile,
+  replacement: any,
+) {
+  // Track inserted nodes and clone them if they are inserted more than
+  // once to avoid injecting the same node multiple times.
+  if (placeholder.isDuplicate) {
+    if (Array.isArray(replacement)) {
+      replacement = replacement.map(node => t.cloneDeep(node));
+    } else if (typeof replacement === "object") {
+      replacement = t.cloneDeep(replacement);
+    }
+  }
+
+  const { parent, key, index } = placeholder.resolve(ast);
+
+  if (placeholder.type === "string") {
+    if (typeof replacement === "string") {
+      replacement = t.stringLiteral(replacement);
+    }
+    if (!replacement || !t.isStringLiteral(replacement)) {
+      throw new Error("Expected string substitution");
+    }
+  } else if (placeholder.type === "statement") {
+    if (index === undefined) {
+      if (!replacement) {
+        replacement = t.emptyStatement();
+      } else if (Array.isArray(replacement)) {
+        replacement = t.blockStatement(replacement);
+      } else if (typeof replacement === "string") {
+        replacement = t.expressionStatement(t.identifier(replacement));
+      } else if (!t.isStatement(replacement)) {
+        replacement = t.expressionStatement((replacement: any));
+      }
+    } else {
+      if (replacement && !Array.isArray(replacement)) {
+        if (typeof replacement === "string") {
+          replacement = t.identifier(replacement);
+        }
+        if (!t.isStatement(replacement)) {
+          replacement = t.expressionStatement((replacement: any));
+        }
+      }
+    }
+  } else if (placeholder.type === "param") {
+    if (typeof replacement === "string") {
+      replacement = t.identifier(replacement);
+    }
+
+    if (index === undefined) throw new Error("Assertion failure.");
+  } else {
+    if (typeof replacement === "string") {
+      replacement = t.identifier(replacement);
+    }
+    if (Array.isArray(replacement)) {
+      throw new Error("Cannot replace single expression with an array.");
+    }
+  }
+
+  if (index === undefined) {
+    t.validate(parent, key, replacement);
+
+    (parent: any)[key] = replacement;
+  } else {
+    const items: Array<BabelNode> = (parent: any)[key].slice();
+
+    if (placeholder.type === "statement" || placeholder.type === "param") {
+      if (replacement == null) {
+        items.splice(index, 1);
+      } else if (Array.isArray(replacement)) {
+        items.splice(index, 1, ...replacement);
+      } else {
+        items[index] = replacement;
+      }
+    } else {
+      items[index] = replacement;
+    }
+
+    t.validate(parent, key, items);
+    (parent: any)[key] = items;
+  }
+}

packages/babel-template/src/string.js

@@ -1,122 +1,23 @@
-import cloneDeep from "lodash/cloneDeep";
-import has from "lodash/has";
-import traverse from "@babel/traverse";
-import * as babylon from "babylon";
-import { codeFrameColumns } from "@babel/code-frame";
+// @flow
+import type { Formatter } from "./formatters";
+import { normalizeReplacements, type TemplateOpts } from "./options";
+import parseAndBuildMetadata from "./parse";
+import populatePlaceholders from "./populate";
 
-const FROM_TEMPLATE = new Set();
+export default function stringTemplate<T>(
+  formatter: Formatter<T>,
+  code: string,
+  opts: TemplateOpts,
+): mixed => T {
+  code = formatter.code(code);
 
-export default function factory(code: string, opts?: Object): Function {
-  // since we lazy parse the template, we get the current stack so we have the
-  // original stack to append if it errors when parsing
-  let stack;
-  try {
-    // error stack gets populated in IE only on throw
-    // (https://msdn.microsoft.com/en-us/library/hh699850(v=vs.94).aspx)
-    throw new Error();
-  } catch (error) {
-    if (error.stack) {
-      // error.stack does not exists in IE <= 9
-      stack = error.stack
-        .split("\n")
-        .slice(2)
-        .join("\n");
-    }
-  }
+  let metadata;
 
-  opts = Object.assign(
-    {
-      allowReturnOutsideFunction: true,
-      allowSuperOutsideMethod: true,
-      preserveComments: false,
-      sourceType: "module",
-    },
-    opts,
-  );
+  return (arg?: mixed) => {
+    const replacements = normalizeReplacements(arg);
 
-  let getAst = function() {
-    let ast;
+    if (!metadata) metadata = parseAndBuildMetadata(formatter, code, opts);
 
-    try {
-      ast = babylon.parse(code, opts);
-
-      ast = traverse.removeProperties(ast, {
-        preserveComments: opts.preserveComments,
-      });
-    } catch (err) {
-      const loc = err.loc;
-      if (loc) {
-        err.loc = null;
-        err.message += "\n" + codeFrameColumns(code, { start: loc });
-      }
-      err.stack = `${err.stack}\n    ==========================\n${stack}`;
-      throw err;
-    }
-
-    getAst = function() {
-      return ast;
-    };
-
-    return ast;
-  };
-
-  return function(...args) {
-    return useTemplate(getAst(), args);
+    return formatter.unwrap(populatePlaceholders(metadata, replacements));
   };
 }
-
-function useTemplate(ast, nodes?: Array<Object>) {
-  ast = cloneDeep(ast);
-  const { program } = ast;
-
-  if (nodes.length) {
-    traverse.cheap(ast, function(node) {
-      FROM_TEMPLATE.add(node);
-    });
-
-    traverse(ast, templateVisitor, null, nodes);
-
-    FROM_TEMPLATE.clear();
-  }
-
-  if (program.body.length > 1) {
-    return program.body;
-  } else {
-    return program.body[0];
-  }
-}
-
-const templateVisitor = {
-  // 360
-  noScope: true,
-
-  Identifier(path, args) {
-    const { node, parentPath } = path;
-    if (!FROM_TEMPLATE.has(node)) return path.skip();
-
-    let replacement;
-    if (has(args[0], node.name)) {
-      replacement = args[0][node.name];
-    } else if (node.name[0] === "$") {
-      const i = +node.name.slice(1);
-      if (args[i]) replacement = args[i];
-    }
-
-    if (parentPath.isExpressionStatement()) {
-      path = parentPath;
-    }
-
-    if (replacement === null) {
-      path.remove();
-    } else if (replacement) {
-      path.replaceInline(replacement);
-      path.skip();
-    }
-  },
-
-  exit({ node }) {
-    if (!node.loc) {
-      traverse.clearNode(node);
-    }
-  },
-};