KaTeX · k4b7 · Dec 22, 2017 · Dec 16, 2017 · Dec 19, 2017 · Dec 20, 2017
diff --git a/src/MacroExpander.js b/src/MacroExpander.js
@@ -7,6 +7,7 @@
 import Lexer, {controlWordRegex} from "./Lexer";
 import {Token} from "./Token";
 import builtinMacros from "./macros";
+import type {Mode} from "./types";
 import ParseError from "./ParseError";
 import objectAssign from "object-assign";
 
@@ -16,13 +17,22 @@ export default class MacroExpander implements MacroContextInterface {
     lexer: Lexer;
     macros: MacroMap;
     stack: Token[];
+    mode: Mode;
 
-    constructor(input: string, macros: MacroMap) {
+    constructor(input: string, macros: MacroMap, mode: Mode) {
         this.lexer = new Lexer(input);
         this.macros = objectAssign({}, builtinMacros, macros);
+        this.mode = mode;
         this.stack = []; // contains tokens in REVERSE order
     }
 
+    /**
+     * Switches between "text" and "math" modes.
+     */
+    switchMode(newMode: Mode) {
+        this.mode = newMode;
+    }
+
     /**
      * Returns the topmost token on the stack, without expanding it.
      * Similar in behavior to TeX's `\futurelet`.

diff --git a/src/Parser.js b/src/Parser.js
@@ -102,9 +102,11 @@ function assertFuncOrArg(parsed) {
 
 export default class Parser {
     constructor(input, settings) {
+        // Start in math mode
+        this.mode = "math";
         // Create a new macro expander (gullet) and (indirectly via that) also a
         // new lexer (mouth) for this parser (stomach, in the language of TeX)
-        this.gullet = new MacroExpander(input, settings.macros);
+        this.gullet = new MacroExpander(input, settings.macros, this.mode);
         // Use old \color behavior (same as LaTeX's \textcolor) if requested.
         // We do this after the macros object has been copied by MacroExpander.
         if (settings.colorIsTextColor) {
@@ -148,6 +150,7 @@ export default class Parser {
      */
     switchMode(newMode) {
         this.mode = newMode;
+        this.gullet.switchMode(newMode);
     }
 
     /**
@@ -157,7 +160,6 @@ export default class Parser {
      */
     parse() {
         // Try to parse the input
-        this.mode = "math";
         this.consume();
         const parse = this.parseInput();
         return parse;
@@ -586,12 +588,16 @@ export default class Parser {
             if (this.mode === "math") {
                 throw new ParseError("$ within math mode");
             }
-            this.consume();
             const outerMode = this.mode;
             this.switchMode("math");
+            // Expand next symbol now that we're in math mode.
+            this.consume();
             const body = this.parseExpression(false, "$");
-            this.expect("$", true);
+            // We can't expand the next symbol after the $ until after
+            // switching modes back.  So don't consume within expect.
+            this.expect("$", false);
             this.switchMode(outerMode);
+            this.consume();
             return new ParseNode("styling", {
                 style: "text",
                 value: body,
@@ -746,29 +752,25 @@ export default class Parser {
      *
      * @return {?ParsedFuncOrArgOrDollar}
      */
-    parseGroupOfType(innerMode, optional) {
-        const outerMode = this.mode;
+    parseGroupOfType(type, optional) {
         // Handle `original` argTypes
-        if (innerMode === "original") {
-            innerMode = outerMode;
+        if (type === "original") {
+            type = this.mode;
         }
 
-        if (innerMode === "color") {
+        if (type === "color") {
             return this.parseColorGroup(optional);
         }
-        if (innerMode === "size") {
+        if (type === "size") {
             return this.parseSizeGroup(optional);
         }
-        if (innerMode === "url") {
+        if (type === "url") {
             return this.parseUrlGroup(optional);
         }
 
-        // By the time we get here, innerMode is one of "text" or "math".
-        // We switch the mode of the parser, recurse, then restore the old mode.
-        this.switchMode(innerMode);
-        const res = this.parseGroup(optional);
-        this.switchMode(outerMode);
-        return res;
+        // By the time we get here, type is one of "text" or "math".
+        // Specify this as mode to parseGroup.
+        return this.parseGroup(optional, type);
     }
 
     consumeSpaces() {
@@ -947,27 +949,38 @@ export default class Parser {
     }
 
     /**
-     * If the argument is false or absent, this parses an ordinary group,
+     * If `optional` is false or absent, this parses an ordinary group,
      * which is either a single nucleus (like "x") or an expression
      * in braces (like "{x+y}").
-     * If the argument is true, it parses either a bracket-delimited expression
+     * If `optional` is true, it parses either a bracket-delimited expression
      * (like "[x+y]") or returns null to indicate the absence of a
      * bracket-enclosed group.
+     * If `mode` is present, switches to that mode while parsing the group,
+     * and switches back after.
      *
      * @param {boolean=} optional  Whether the group is optional or required
      * @return {?ParsedFuncOrArgOrDollar}
      */
-    parseGroup(optional) {
+    parseGroup(optional, mode) {
+        const outerMode = this.mode;
         const firstToken = this.nextToken;
         // Try to parse an open brace
         if (this.nextToken.text === (optional ? "[" : "{")) {
+            // Switch to specified mode before we expand symbol after brace
+            if (mode) {
+                this.switchMode(mode);
+            }
             // If we get a brace, parse an expression
             this.consume();
             const expression = this.parseExpression(false, optional ? "]" : "}");
             const lastToken = this.nextToken;
+            // Switch mode back before consuming symbol after close brace
+            if (mode) {
+                this.switchMode(outerMode);
+            }
             // Make sure we get a close brace
             this.expect(optional ? "]" : "}");
-            if (this.mode === "text") {
+            if (mode === "text") {
                 this.formLigatures(expression);
             }
             return newArgument(
@@ -976,7 +989,14 @@ export default class Parser {
                     firstToken.range(lastToken, firstToken.text));
         } else {
             // Otherwise, just return a nucleus, or nothing for an optional group
-            return optional ? null : this.parseSymbol();
+            if (mode) {
+                this.switchMode(mode);
+            }
+            const result = optional ? null : this.parseSymbol();
+            if (mode) {
+                this.switchMode(outerMode);
+            }
+            return result;
         }
     }
 

diff --git a/src/macros.js b/src/macros.js
@@ -52,11 +52,23 @@ function defineMacro(name: string, body: MacroDefinition) {
 //////////////////////////////////////////////////////////////////////
 // macro tools
 
+// LaTeX's \@firstoftwo{#1}{#2} expands to #1, skipping #2
+// TeX source: \long\def\@firstoftwo#1#2{#1}
 defineMacro("\\@firstoftwo", function(context) {
     const args = context.consumeArgs(2);
     return {tokens: args[0], numArgs: 0};
 });
 
+// LaTeX's \@secondoftwo{#1}{#2} expands to #2, skipping #1
+// TeX source: \long\def\@secondoftwo#1#2{#2}
+defineMacro("\\@secondoftwo", function(context) {
+    const args = context.consumeArgs(2);
+    return {tokens: args[1], numArgs: 0};
+});
+
+// LaTeX's \@ifnextchar{#1}{#2}{#3} looks ahead to the next (unexpanded)
+// symbol.  If it matches #1, then the macro expands to #2; otherwise, #3.
+// Note, however, that it does not consume the next symbol in either case.
 defineMacro("\\@ifnextchar", function(context) {
     const args = context.consumeArgs(3);  // symbol, if, else
     const nextToken = context.future();
@@ -67,9 +79,22 @@ defineMacro("\\@ifnextchar", function(context) {
     }
 });
 
-// \def\@ifstar#1{\@ifnextchar *{\@firstoftwo{#1}}}
+// LaTeX's \@ifstar{#1}{#2} looks ahead to the next (unexpanded) symbol.
+// If it is `*`, then it consumes the symbol, and the macro expands to #1;
+// otherwise, the macro expands to #2 (without consuming the symbol).
+// TeX source: \def\@ifstar#1{\@ifnextchar *{\@firstoftwo{#1}}}
 defineMacro("\\@ifstar", "\\@ifnextchar *{\\@firstoftwo{#1}}");
 
+// LaTeX's \TextOrMath{#1}{#2} expands to #1 in text mode, #2 in math mode
+defineMacro("\\TextOrMath", function(context) {
+    const args = context.consumeArgs(2);
+    if (context.mode === 'text') {
+        return {tokens: args[0], numArgs: 0};
+    } else {
+        return {tokens: args[1], numArgs: 0};
+    }
+});
+
 //////////////////////////////////////////////////////////////////////
 // basics
 defineMacro("\\bgroup", "{");

diff --git a/test/katex-spec.js b/test/katex-spec.js
@@ -2672,11 +2672,61 @@ describe("A macro expander", function() {
         expect("\\@ifnextchar!{yes}{no}?!").toParseLike("no?!");
     });
 
-    it("\\@firstoftwwo should consume star but nothing else", function() {
+    it("\\@ifstar should consume star but nothing else", function() {
         expect("\\@ifstar{yes}{no}*!").toParseLike("yes!");
         expect("\\@ifstar{yes}{no}?!").toParseLike("no?!");
     });
 
+    it("\\TextOrMath should work immediately", function() {
+        expect("\\TextOrMath{text}{math}").toParseLike("math");
+    });
+
+    it("\\TextOrMath should work after other math", function() {
+        expect("x+\\TextOrMath{text}{math}").toParseLike("x+math");
+    });
+
+    it("\\TextOrMath should work immediately after \\text", function() {
+        expect("\\text{\\TextOrMath{text}{math}}").toParseLike("\\text{text}");
+    });
+
+    it("\\TextOrMath should work later after \\text", function() {
+        expect("\\text{hello \\TextOrMath{text}{math}}")
+            .toParseLike("\\text{hello text}");
+    });
+
+    it("\\TextOrMath should work immediately after \\text ends", function() {
+        expect("\\text{\\TextOrMath{text}{math}}\\TextOrMath{text}{math}")
+            .toParseLike("\\text{text}math");
+    });
+
+    it("\\TextOrMath should work immediately after $", function() {
+        expect("\\text{$\\TextOrMath{text}{math}$}")
+            .toParseLike("\\text{$math$}");
+    });
+
+    it("\\TextOrMath should work later after $", function() {
+        expect("\\text{$x+\\TextOrMath{text}{math}$}")
+            .toParseLike("\\text{$x+math$}");
+    });
+
+    it("\\TextOrMath should work immediately after $ ends", function() {
+        expect("\\text{$\\TextOrMath{text}{math}$\\TextOrMath{text}{math}}")
+            .toParseLike("\\text{$math$text}");
+    });
+
+    it("\\TextOrMath should work in a macro", function() {
+        compareParseTree("\\mode\\text{\\mode$\\mode$\\mode}\\mode",
+            "math\\text{text$math$text}math",
+            {"\\mode": "\\TextOrMath{text}{math}"});
+    });
+
+    // TODO(edemaine): This doesn't work yet.  Parses like `\text math`,
+    // which doesn't even treat all four letters as an argument.
+    //it("\\TextOrMath should work in a macro passed to \\text", function() {
+    //    compareParseTree("\\text\\mode", "\\text{text}",
+    //        {"\\mode": "\\TextOrMath{text}{math}"});
+    //});
+
     // This may change in the future, if we support the extra features of
     // \hspace.
     it("should treat \\hspace, \\hspace*, \\hskip like \\kern", function() {