Optimize.cpp

/* Attempts to reduces the number of lines of code in a TM by looking
 * for patterns commonly produced by the code generators for Call/Done
 * and for Set/Load.
 *
 * This optimizer works as follows.
 *
 *   * For each label consisting just of an unconditional Goto, replace all
 *     Gotos to that label with Gotos to the target label.
 *
 *   * Remove all (non-Start) labels that have nothing pointing into them.
 *
 *   * Reorder the code so that Gotos that are the only way to reach
 *     specific labels are removed, along with the target label.
 *
 *   * Remove all Gotos that go directly to the next line.
 *
 *   * Group all consecutive labels together into a single label.
 *
 *   * Remove any Move Lefts immediately followed by Remove Rights
 *     or vice-versa.
 *
 *   * Repeat the above until no improvements can be made.
 *
 *   * Replace all label names with (hopefully!) shorter ones.
 *
 * These are all syntactic simplifications; no attempt is made at
 * semantic improvements (e.g. looking for unnecessary moves / writes
 * or combining different code paths together).
 */
#include <iostream>
#include <string>
#include <fstream>
#include <vector>
#include <iomanip>
#include <queue>
#include <set>
#include <sstream>
#include <optional>
#include "Turing/StrUtils/StrUtils.h"
#include "Turing/Turing.h"
using namespace std;

namespace {
    /* Lists all errors found in the program. */
    void showErrorsIn(const Turing::Program& program) {
        cout << "Errors found in program:" << endl;
        for (size_t i = 0; i < program.numLines(); i++) {
            auto error = program.errorAtLine(i);
            if (error != "") {
                cout << setw(6) << i << " " << program.line(i) << endl;
                cout << "           " << error << endl;
            }
        }
    }
    
    /* Finds all Goto labels. */
    set<string> usedLabelsIn(const Turing::Program& program) {
        set<string> result = { "Start" }; // Start is always used
    
        for (size_t i = 0; i < program.numLines(); i++) {
            if (auto statement = program.statement(i)) {
                class LabelFinder: public Turing::Calculator<set<string>> {
                public:
                    set<string> handle(Turing::Goto& g) override {
                        return { g.label() };
                    }
                    set<string> handle(Turing::If& i, const set<string>& nested) override {
                        return nested;
                    }
                };
                
                auto labels = LabelFinder().calculate(*statement);
                result.insert(labels.begin(), labels.end());
            }
        }
        
        return result;
    }
    
    Turing::Program removeUnusedLabelsFrom(const Turing::Program& p) {
        auto labels = usedLabelsIn(p);
        
        stringstream builder;
        for (size_t i = 0; i < p.numLines(); i++) {
            auto label = dynamic_pointer_cast<Turing::Label>(p.statement(i));
            
            /* Output this line UNLESS it's an unused label (i.e. not a label or
             * is used) OR has no statement.
             */
            if (p.statement(i) && (!label || labels.contains(label->label()))) {
                builder << p.line(i) << '\n';
            }
        }
        
        return Turing::Program(builder);
    }
    
    /* A basic block in the program. This consists of a maximal region of
     * contiguous straight-line code.
     */
    struct BasicBlock {
        /* Range of lines from the original program. */
        size_t firstLine = 0xBADBADBADU, lastLine = 0xBADBADBADU;
        
        /* In and out arrows in the control-flow graph. */
        set<BasicBlock*> in, out;
        
        /* "Pinning" information. A basic block can be pinned
         * to the block before or after it if it continues to
         * those blocks via fallthrough. When moving blocks
         * around, we have to treat pinned blocks as a unit.
         *
         * Each block stores the BB that is pinned before
         * and after it, which could be nullptr if it isn't
         * pinned on each side.
         */
        BasicBlock* pinnedBefore = nullptr;
        BasicBlock* pinnedAfter  = nullptr;
        
        /* Whether this basic block terminates in an unconditional
         * Goto statement. This is used by the optimizer to reorder
         * blocks and eliminate Gotos and labels.
         */
        bool endsInGoto = false;
        
        /* Whether this basic block begins with a label. This is used
         * by the optimizer to reorder blocks and eliminate Gotos and
         * labels.
         */
        bool startsWithLabel = false;
        bool isStartLabel    = false;
        
        /* The "chain" that this basic block belongs to. */
        struct Chain* chain = nullptr;
    };
    
    /* A "chain" in the program. A chain consists of a maximal region
     * of pinned blocks in a row.
     */
    struct Chain {
        deque<BasicBlock*> bbs;
    };
    
    /* Returns whether the given statement is a label. */
    bool isLabel(const shared_ptr<Turing::Statement>& statement) {
        return dynamic_pointer_cast<Turing::Label>(statement) != nullptr;
    }
    
    /* Returns whether the given statement contains a Goto. */
    bool containsGoto(shared_ptr<Turing::Statement> statement) {
        /* Peek into If statements. */
        if (auto i = dynamic_pointer_cast<Turing::If>(statement)) {
            statement = i->stmt();
        }
        return dynamic_pointer_cast<Turing::Goto>(statement) != nullptr;
    }
    
    /* Returns whether the given statement contains a Return. */
    bool containsReturn(shared_ptr<Turing::Statement> statement) {
        /* Peek into If statements. */
        if (auto i = dynamic_pointer_cast<Turing::If>(statement)) {
            statement = i->stmt();
        }
        return dynamic_pointer_cast<Turing::Return>(statement) != nullptr;
    }
    
    /* Given a line in a program, outputs the line numbers of its successors. */
    set<size_t> successorsOf(const Turing::Program& program, size_t line) {
        class SuccessorFinder: public Turing::Calculator<set<size_t>> {
        public:
            SuccessorFinder(const Turing::Program& program, size_t line) : program(program), line(line) {
                // Handled in initialization line
            }
        
            /* Most statements fall through to the next line. */
            set<size_t> handle(Turing::Label&) override {
                return nextLine();
            }
            set<size_t> handle(Turing::Move&) override {
                return nextLine();
            }
            set<size_t> handle(Turing::Write&) override {
                return nextLine();
            }
            
            /* Return statements have no successors. */
            set<size_t> handle(Turing::Return&) override {
                return { };
            }
            
            /* Conditionals can either go to the next line or wherever their
             * nested statements go.
             */
            set<size_t> handle(Turing::If&, const set<size_t>& nestedTargets) override {
                auto result = nestedTargets;
                auto next   = nextLine();
                result.insert(next.begin(), next.end());
                return result;
            }
            
            /* Gotos go where their labels are. */
            set<size_t> handle(Turing::Goto& g) override {
                return { program.lineForLabel(g.label()) };
            }
        
        private:
            const Turing::Program& program;
            size_t line;
            
            /* Finds the next line after the given one. */
            set<size_t> nextLine() const {
                return line + 1 < program.numLines()? set<size_t>{ line + 1 } : set<size_t>{ };
            }
        };
        
        return SuccessorFinder(program, line).calculate(*program.statement(line));
    }
    
    /* Basic blocks, keyed by their starting lines. */
    map<size_t, BasicBlock*> basicBlocksIn(const Turing::Program& program) {
        map<size_t, BasicBlock*> result;
        
        /* Worklist algorithm; we're doing a BFS over the program looking
         * for reachability. As we do, we assemble the basic blocks.
         */
        auto startLine = program.lineForLabel("Start");         
        queue<size_t> worklist;
        worklist.push(startLine);
        result[startLine] = new BasicBlock { startLine, startLine };
        
        /* Standard BFS. */
        while (!worklist.empty()) {
            /* Grab the current line. */
            size_t startLine = worklist.front();
            worklist.pop();
            
            /* Keep stepping forward until we have a reason not to. */
            auto* bb = result[startLine];
            for (size_t line = startLine; line < program.numLines(); line++) {
                /* If this line is a label and isn't the very first line in the
                 * BB, we've gone too far.
                 */
                if (line != startLine && isLabel(program.statement(line))) {
                    break;
                }
            
                /* Otherwise, this line is part of the BB. Pessimistically assume
                 * it's the last line.
                 */
                bb->lastLine = line;
                
                /* The basic block will continue onward to the next line unless any of
                 * the following conditions hold:
                 *
                 *   1. The current line is a Goto or contains a Goto.
                 *   2. The current line is a Return or contains a Return.
                 */
                if (containsGoto(program.statement(line)) ||
                    containsReturn(program.statement(line))) {
                    break;
                }
            }
            
            /* Now, wire the appropriate links to other blocks. */
            for (size_t successor: successorsOf(program, bb->lastLine)) {
                auto* targetBB = result[successor];
                
                /* If we've never seen this BB before, we need to enqueue it into our BFS
                 * worklist. If we have seen it before, all we're doing is updating its
                 * information.
                 */
                if (!targetBB) {
                    result[successor] = new BasicBlock{ successor, successor };
                    targetBB = result[successor];
                    worklist.push(successor);
                }
                
                /* Update predecessor/successor information. */
                targetBB->in.insert(bb);
                bb->out.insert(targetBB);
            }
            
            /* Look at the very last instruction of the current basic block. If it's
             * not a Goto or Return, then we are pinned to whatever block comes after us.
             *
             * Edge case: The very last line of the program can't pin to the nothing
             * that comes after it.
             */
            if (typeid(*program.statement(bb->lastLine)) != typeid(Turing::Goto) &&
                typeid(*program.statement(bb->lastLine)) != typeid(Turing::Return) &&
                bb->lastLine + 1 != program.numLines()) {
                
                auto* successor = result.at(bb->lastLine + 1);
                bb->pinnedAfter = successor;
                successor->pinnedBefore = bb;
            }
            
            /* If our very last line is an unconditional Goto, remember that for later. */
            if (typeid(*program.statement(bb->lastLine)) == typeid(Turing::Goto)) {
                bb->endsInGoto = true;
            }
            
            /* If our very first line is a label, remember that for later. */
            if (typeid(*program.statement(bb->firstLine)) == typeid(Turing::Label)) {
                bb->startsWithLabel = true;
            }
            
            /* If we're the start label, remember that for later. */
            if (program.lineForLabel("Start") == bb->firstLine) {
                bb->isStartLabel = true;
            }
        }
        
        return result;
    }
    
    [[ maybe_unused ]] void printBBs(const map<size_t, BasicBlock*>& bbs, const Turing::Program& program) {
        for (auto [line, bb]: bbs) {
            cout << "Basic block starting on line " << line << ": " << endl;
            
            for (size_t i = bb->firstLine; i <= bb->lastLine; i++) {
                cout << "    " << setw(6) << i << " " << program.line(i) << endl;
            }
            
            cout << "    ----" << endl;
            cout << "    IN: ";
            for (auto* in: bb->in) {
                cout << in->firstLine << " ";
            }
            cout << endl;
            cout << "    OUT: ";
            for (auto* out: bb->out) {
                cout << out->firstLine << " ";
            }
            cout << endl;
            if (bb->pinnedBefore) {
                cout << "    Pinned above to BB " << bb->pinnedBefore->firstLine << endl;
            }
            if (bb->pinnedAfter) {
                cout << "    Pinned below to BB " << bb->pinnedAfter->firstLine << endl;
            }
            if (bb->endsInGoto) {
                cout << "    Block ends with Goto." << endl;
            }
            if (bb->startsWithLabel) {
                cout << "    Block starts with a label." << endl;
            }
            
            cout << endl;
        }
    }
    
    /* Given a basic block, returns the basic block at the top of its pin chain. */
    BasicBlock* topOfPinChain(BasicBlock* bb) {
        while (bb->pinnedBefore) bb = bb->pinnedBefore;
        return bb;
    }
    
    /* Utility to get a nice printed version of a BB's line range. */
    [[ maybe_unused ]] string spanOf(BasicBlock* bb) {
        ostringstream builder;
        builder << bb->firstLine << " - "
                << bb->lastLine;
        return builder.str();
    }
    /* Utility to get a nice printed version of a chain's line range. */
    [[ maybe_unused ]] string spanOf(Chain* chain) {
        ostringstream builder;
        builder << chain->bbs.front()->firstLine << " - "
                << chain->bbs.back()->lastLine;
        return builder.str();
    }
    
    /* Computes the chains associated with each BB by following up and down the
     * pin links.
     */
    void determineChainsIn(const map<size_t, BasicBlock*>& bbs, const Turing::Program& program) {
        /* Standard BFS. We use the first entry in the chain as the element of the 
         * visited set.
         */
        queue<BasicBlock*> worklist;
        worklist.push(topOfPinChain(bbs.at(program.lineForLabel("Start"))));
        set<BasicBlock*> visited = { worklist.front() };
        
        while (!worklist.empty()) {
            auto* bb = worklist.front();
            worklist.pop();
            
            /* Form the chain one link at a time. */
            auto* chain = new Chain;
            for (auto* curr = bb; curr != nullptr; curr = curr->pinnedAfter) {
                curr->chain = chain;
                chain->bbs.push_back(curr);
                
                /* Add successors to the worklist. */
                for (auto* successor: curr->out) {
                    auto* top = topOfPinChain(successor);
                    
                    if (!visited.contains(top)) {
                        visited.insert(top);
                        worklist.push(top);
                    }
                }
            }
        }
    }
    
    [[ maybe_unused ]] void printChains(const map<size_t, BasicBlock*>& bbs,
                                        const Turing::Program& program) {
        set<Chain*> used;
        for (auto [ _, bb ]: bbs) {
            if (!used.contains(bb->chain)) {
                used.insert(bb->chain);
                
                cout << "Next chain: " << endl;
                
                for (size_t i = bb->chain->bbs.front()->firstLine; i <= bb->chain->bbs.back()->lastLine; i++) {
                    cout << setw(6) << i << " " << program.line(i) << endl;
                }
            }
        }                                    
    }
    
    /* Returns the predecessor chain, if any, in the context of the Goto/Label elimination
     * optimization. The predecessor, if one exists, is the chain that ends with a Goto,
     * has indegree one, and jumps to the label at the top of this chain.
     */
    Chain* predecessorOf(Chain* chain) {    
        /* Look at the basic block at the front of the chain. */
        auto* firstBB = chain->bbs.front();
        if (firstBB->startsWithLabel && firstBB->in.size() == 1) {
            /* See which BB points at us. */
            auto* inBB = *firstBB->in.begin();
            
            /* If the BB ends in a Goto, it's also definitionally the
             * end of a chain because there's no fallthrough.
             */
            if (inBB->endsInGoto) return inBB->chain;
        }
        
        /* Nope, has no predecessor. */
        return nullptr;
    }
    
    /* Returns the successor chain, if any, in the context of the Goto/Label elimination
     * optimization. The successor, if one exists, is the chain that starts with a label,
     * has indegree 1, and is pointed at by an unconditional Goto at the end of this
     * chain.
     */
    Chain* successorOf(Chain* chain) {
        /* Look at the last BB in the chain. */
        auto* lastBB = chain->bbs.back();
        if (lastBB->endsInGoto) {
            /* Now see where we go. */
            auto* outBB = *lastBB->out.begin();
            
            /* If it has indegree one, it's our successor. */
            if (outBB->in.size() == 1) return outBB->chain;
        }
        
        /* Nope, no successor. */
        return nullptr;
    }
    
    /* A major optimization pass that reorders the code to put labels only reached
     * by a single Goto right after that Goto. A side-effect of this optimization
     * is that we eliminate all unreachable code.
     */
    Turing::Program reorderCode(const Turing::Program& program) {
        stringstream builder;    
        
        /* Find the basic blocks in this program. */
        auto bbs = basicBlocksIn(program);
        
        /* Determine the chains for each (reachable) basic block. */
        determineChainsIn(bbs, program);
        
        /* We can visit the basic blocks in any order as long as we don't repeat any
         * chains. The BB we initiate from is not necessarily going to be the first BB
         * that's emitted.
         */
        set<Chain*> used;
        for (auto [_, bb]: bbs) {    
            /* See if we need to process this chain. */
            auto* chain = bb->chain;
            if (!chain) {
                continue;
            }
            
            if (used.contains(chain)) {
                continue;
            }
            
            /* Here's how the optimization works. We are looking for chains where
             * the end of one chain is an unconditional Goto and the start of the
             * destination chain is a label that has indegree one. This means that
             * the Goto from the destination chain, and the label, can (likely) be
             * eliminated by reordering the chains so that the second chain starts
             * where the first one ends. We then drop the label and the Goto, saving
             * two lines.
             *
             * The graph structure of the chains is a digraph where each node has
             * indegree at most one and outdegree at most one. Why? Because
             *
             *   INDEGREE AT MOST ONE: Each edge corresponds to something that
             *   occurs when the start of the basic block is a label of indegree
             *   one. Therefore, there can be at most one incoming edge.
             *
             *   OUTDEGREE AT MOST ONE: Each edge originates at an unconditional
             *   Goto, which can only go to one spot.
             *
             * This partitions the chains into paths and cycles. Paths are easy to
             * output: you just start at the beginning of the path and lay things
             * out until you get to the end. Cycles are a bit trickier because the
             * naive approach would run into an infinite loop trying to lay out
             * blocks before other blocks. (Ask me how I learned that! >_<)
             *
             * Thus in the cycle case, we simply cut the cycle at some arbitrary
             * point and leave the very first label and very last Goto unmodified.
             *
             * We also need to be careful not to accidentally remove the Start
             * label. That one has to remain there.
             */
             
            /* Order of the chains that need to be emitted. */
            deque<Chain*> toOutput = { chain };
            bool hasLoop = false;
            
            /* Walk backwards from this point. We will either hit a cycle or
             * we will find a node of indegree 0.
             */
            for (auto* curr = chain; (curr = predecessorOf(curr)); ) {
                /* We'll hit a loop if we return to our starting chain. */
                if (curr == toOutput.back()) {
                    hasLoop = true;
                    break;
                }
            
                /* Otherwise prepend this to our list. */
                toOutput.push_front(curr);
            }
            
            /* If we didn't hit a cycle, walk forwards until we reach a chain
             * with no successor.
             */
            if (!hasLoop) for (auto* curr = chain; (curr = successorOf(curr)); ) {
                toOutput.push_back(curr);
            }
            
            /* Now, render these from top to bottom. If there's a loop, don't
             * apply the optimization to the first label and last Goto.
             */
            for (auto* chain: toOutput) {
                for (auto* bb: chain->bbs) {
                    size_t start = bb->firstLine;
                    
                    /* Retain the first label in the chain, as we still need that
                     * as a Goto target. (Otherwise, we would have something before us.)
                     * Otherwise skip the label at the start of the chain.
                     *
                     * Oh... and never remove Start!
                     */
                    if (bb == chain->bbs.front() && chain != toOutput.front() && !bb->isStartLabel) {
                        start++;
                    }
                    
                    size_t end = bb->lastLine;
                    
                    /* Retain the last Goto in the chain, as otherwise the chain
                     * would have continued.
                     */
                    if (bb == chain->bbs.back() && chain != toOutput.back()) {
                        end--;
                    }
                    
                    /* Output the line. */
                    for (size_t line = start; line <= end; line++) {
                        builder << program.line(line) << '\n';
                    }
                }
                
                /* Remember that we output this so that we don't do it again. */
                used.insert(chain);
            }
        }
        
        return Turing::Program(builder);
    }
    
    /* Remove any Gotos that jump to the next line. */
    Turing::Program removeUnneededGotos(const Turing::Program& program) {
        stringstream builder;
        
        /* Remove any line that's a Goto that points to the next line. */
        for (size_t i = 0; i < program.numLines(); i++) {
            bool writeLine = true;
            if (auto g = dynamic_pointer_cast<Turing::Goto>(program.statement(i))) {
                /* No need for a Goto that just goes to the next line. */
                if (program.lineForLabel(g->label()) == i + 1) {
                    writeLine = false;
                }
            }
            
            if (writeLine) builder << program.line(i) << '\n';
        }
        
        return Turing::Program(builder);
    }
    
    /* Calculator that rewrites programs by deleting labels whose new names don't
     * match their old names and that rewrites Goto targets.
     */
    class Rewriter: public Turing::Calculator<Turing::Statement*> {
    public:
        Rewriter(const map<string, string>& newNames): newNames(newNames) {
        
        }
        
        /* Moves are unchanged. */
        Turing::Statement* handle(Turing::Move& m) override {
            return &m;
        }
        
        /* Writes are unchanged. */
        Turing::Statement* handle(Turing::Write& w) override {
            return &w;
        }
        
        /* Returns are unmodified. */
        Turing::Statement* handle(Turing::Return& r) override {
            return &r;
        }
        
        /* Gotos use the new target. */
        Turing::Statement* handle(Turing::Goto& g) override {
            return new Turing::Goto(newNames.at(g.label()));
        }
        
        /* Labels that map to new names get dropped; labels that keep
         * their name are retained.
         */
        Turing::Statement* handle(Turing::Label& l) override {
            return newNames.at(l.label()) == l.label()? &l : nullptr;
        }
        
        /* If statements use the translated nested statement. */
        Turing::Statement* handle(Turing::If& i, Turing::Statement* const & nested) override {
            return new Turing::If(i.isNegated(), i.ch(), shared_ptr<Turing::Statement>(nested));
        }
        
    
    private:
        const map<string, string>& newNames;
    };
    
    /* Symbol to text representation. */
    string symbolToString(char32_t ch) {
        if (ch == Turing::kBlankSymbol) {
            return "Blank";
        } else if (ch == '\n') {
            return "'\\n'";
        } else if (ch == '\t') {
            return "'\\t'";
        } else if (ch == '\'') {
            return "'\\''";
        } else if (ch == '\\') {
            return "'\\\\'";
        } else {
            return "\'" + string(1, ch) + "\'";
        }
    }
    
    /* Calculator that converts statements to strings. */
    class StatementToString: public Turing::Calculator<string> {
    public:
        string handle(Turing::Label& l) override {
            return l.label() + ":";
        }
        string handle(Turing::Move& m) override {
            return "Move " + string(m.direction() == Turing::Direction::LEFT? "Left" : "Right");
        }
        string handle(Turing::Return& r) override {
            return "Return " + string(r.isAccepting()? "True" : "False");
        }
        string handle(Turing::Goto& g) override {
            return "Goto " + g.label();
        }
        string handle(Turing::Write& w) override {
            return "Write " + symbolToString(w.ch());
        }
        string handle(Turing::If& w, const string& nested) override {
            return "If " + string(w.isNegated()? "Not " : "") + symbolToString(w.ch()) + " " + nested;
        }
    };
    
    /* Look for consecutive labels and replace them with one unified label.
     * This requires two passes: one to figure out where those labels are and
     * rename them, and then one to rewrite the labels.
     */
    Turing::Program removeConsecutiveLabels(const Turing::Program& program) {
        map<string, string> newNames; // Old name -> New name
        
        /* Phase 1: Assign new names to the labels. */
        for (size_t i = 0; i < program.numLines(); ) {
            /* If it's a label, scan forward while we're finding more labels. */
            if (auto l = dynamic_pointer_cast<Turing::Label>(program.statement(i))) {
                set<string> groupNames = { l->label() };
                
                /* Find all consecutive labels. */
                size_t lastLabel = i + 1;
                for (; lastLabel < program.numLines(); lastLabel++) {
                    if (auto l2 = dynamic_pointer_cast<Turing::Label>(program.statement(lastLabel))) {
                        groupNames.insert(l2->label());
                    } else break;
                }
                
                /* Pick a unified name for all the labels. If any of them are Start use that; else
                 * just pick the first.
                 */
                string newName = (groupNames.contains("Start")? "Start" : *groupNames.begin());
                for (string name: groupNames) {
                    newNames[name] = newName;
                }
                
                /* Advance forward. */
                i = lastLabel;
            } else i++;
        }
        
        /* Phase 2: Rewrite the program by removing duplicate labels and changing Goto
         * targets.
         */
        stringstream builder;
        for (size_t i = 0; i < program.numLines(); i++) {           
            if (auto* result = Rewriter(newNames).calculate(*program.statement(i))) {
                builder << StatementToString().calculate(*result) << '\n';
            }            
        }
        
        return Turing::Program(builder);
    }
    
    /* Look for consecutive Writes and remove all but the last of them. */
    Turing::Program removeConsecutiveWrites(const Turing::Program& program) {        
        stringstream builder;
        for (size_t i = 0; i < program.numLines(); ) {
            /* If it's a Write, scan forward while we're finding more Writes. */
            if (auto w = dynamic_pointer_cast<Turing::Write>(program.statement(i))) {
                size_t lastWrite = i + 1;
                for (; lastWrite < program.numLines(); lastWrite++) {
                    if (!dynamic_pointer_cast<Turing::Write>(program.statement(lastWrite))) break;
                }
                
                /* Only output the last write. */
                builder << program.line(lastWrite - 1) << '\n';
                                
                /* Advance forward. */
                i = lastWrite;
            }
            /* Write everything else unmodified. */
            else {
                builder << program.line(i) << '\n';
                i++;
            }
        }
        
        return Turing::Program(builder);
    }
    
    /* Command-line args telling us what to do. */
    struct Arguments {
        bool simplifyNames = true;
    };
    
    Arguments parseArguments(int argc, const char* argv[]) {
        /* We either take no args or the arg --keep-names. */
        if (argc == 1) return {};
        if (argc == 2 && argv[1] == string("--keep-names")) return { false };
        
        cout << "Usage: optimize [--keep-names]" << endl;
        exit(-1);
    }
    
    /* Calculator that rewrites programs by renaming labels. */
    class Renamer: public Turing::Calculator<Turing::Statement*> {
    public:
        Renamer(const map<string, string>& newNames): newNames(newNames) {
        
        }
        
        /* Moves are unchanged. */
        Turing::Statement* handle(Turing::Move& m) override {
            return &m;
        }
        
        /* Writes are unchanged. */
        Turing::Statement* handle(Turing::Write& w) override {
            return &w;
        }
        
        /* Returns are unmodified. */
        Turing::Statement* handle(Turing::Return& r) override {
            return &r;
        }
        
        /* Gotos use the new target. */
        Turing::Statement* handle(Turing::Goto& g) override {
            return new Turing::Goto(newNames.at(g.label()));
        }
        
        /* Labels that map to new names get dropped; labels that keep
         * their name are retained.
         */
        Turing::Statement* handle(Turing::Label& l) override {
            return new Turing::Label(newNames.at(l.label()));
        }
        
        /* If statements use the translated nested statement. */
        Turing::Statement* handle(Turing::If& i, Turing::Statement* const & nested) override {
            return new Turing::If(i.isNegated(), i.ch(), shared_ptr<Turing::Statement>(nested));
        }
        
    
    private:
        const map<string, string>& newNames;
    };
    
    /* Calculator that rewrites programs by rewriting Gotos while preserving label names. */
    class GotoTranslator: public Turing::Calculator<Turing::Statement*> {
    public:
        GotoTranslator(const map<string, string>& newNames): newNames(newNames) {
        
        }
        
        /* Moves are unchanged. */
        Turing::Statement* handle(Turing::Move& m) override {
            return &m;
        }
        
        /* Writes are unchanged. */
        Turing::Statement* handle(Turing::Write& w) override {
            return &w;
        }
        
        /* Returns are unmodified. */
        Turing::Statement* handle(Turing::Return& r) override {
            return &r;
        }
        
        /* Gotos use the new target. */
        Turing::Statement* handle(Turing::Goto& g) override {
            return new Turing::Goto(newNames.at(g.label()));
        }
        
        /* Labels that map to new names get dropped; labels that keep
         * their name are retained.
         */
        Turing::Statement* handle(Turing::Label& l) override {
            return &l;
        }
        
        /* If statements use the translated nested statement. */
        Turing::Statement* handle(Turing::If& i, Turing::Statement* const & nested) override {
            return new Turing::If(i.isNegated(), i.ch(), shared_ptr<Turing::Statement>(nested));
        }
        
    
    private:
        const map<string, string>& newNames;
    };
    
    /* Returns the set of all legal characters that can be in a label. */
    string legalLabelValues() {
        static string result = [] {
            string result;
            for (int i = 0; i < 128; i++) {
                if (isalnum(i) || i == '_') {
                    result += i;
                }
            }
            return result;
        } ();
        
        return result;
    }
    
    /* Given a base-63 string, advances to the next one. */
    string nextLabel(string input) {
        string chars = legalLabelValues();
        
        for (size_t i = 0; i < input.size(); i++) {
            /* Increment this digit. If no overflow, we're done. */
            size_t value = chars.find(input[i]);
            if (value + 1 != chars.size()) {
                input[i] = chars[value + 1];
                return input;
            }
            
            /* Otherwise there is overflow. That's fine. */
            input[i] = chars[0];
        }
        
        /* If we're here, we need another digit. */
        return input + chars[0];
    }
    
    /* Replaces all non-Start label names with something shorter. */
    Turing::Program simplifyNames(const Turing::Program& program) {
        map<string, string> newNames = { { "Start", "Start" } };
        
        /* Phase 1: Come up with new names for all the labels. */
                 
        string label(1, legalLabelValues()[0]);
        for (size_t i = 0; i < program.numLines(); i++) {
            if (auto l = dynamic_pointer_cast<Turing::Label>(program.statement(i))) {
                if (l->label() != "Start") {
                    /* Prepend an underscore to each label to ensure it's syntactically
                     * valid.
                     */
                    newNames[l->label()] = "_" + label;
                    label = nextLabel(label);
                }
            }
        }
        
        /* Phase 2: Translate all labels. */
        stringstream builder;
        for (size_t i = 0; i < program.numLines(); i++) {
            builder << StatementToString().calculate(*Renamer(newNames).calculate(*program.statement(i))) << '\n';
        }
        
        return Turing::Program(builder);
    }
    
    /* If we see 'Move Left' followed by 'Move Right' or vice-versa, we should output neither. */
    Turing::Program removeUselessMoves(const Turing::Program& program) {
        stringstream builder;
        
        for (size_t i = 0; i + 1 < program.numLines(); i++) {        
            /* See if this and the next line are moves. */
            auto m1 = dynamic_pointer_cast<Turing::Move>(program.statement(i));            
            auto m2 = dynamic_pointer_cast<Turing::Move>(program.statement(i + 1));
            
            /* If so, see if they go in opposite directions. */
            if (m1 && m2 && m1->direction() != m2->direction()) {
                /* Don't output this line. Also, skip the next line. */
                i++;
            }
            /* Otherwise output as usual. */
            else {
                builder << program.line(i) << '\n';
            }
        }
        
        /* Always output the last line. */
        builder << program.line(program.numLines() - 1) << '\n';
        
        return Turing::Program(builder);
    }
    
    /* Look for labels of the following form:
     *
     *    Label:
     *        Goto OtherLabel
     *
     * Anything that points to Label should instead point to where OtherLabel
     * points. Moreover, this can be chained, so anything jumping to Label
     * should instead go where OtherLabel ultimately points.
     */
    Turing::Program bypassUnnecessaryLabels(const Turing::Program& program) {
        /* Map from labels to their ultimate replacements. Labels that don't
         * meet the criteria here will be replaced by themselves.
         */
        map<string, string> newNames;
        
        for (size_t i = 0; i < program.numLines(); i++) {
            if (auto l = dynamic_pointer_cast<Turing::Label>(program.statement(i))) {
                /* Default to the label replacing itself. */
                newNames[l->label()] = l->label();
            
                /* See if there's a line after us and whether it's a goto. */
                if (i + 1 != program.numLines()) {
                    if (auto g = dynamic_pointer_cast<Turing::Goto>(program.statement(i + 1))) {
                        /* Our new target is that target. */
                        newNames[l->label()] = g->label();
                    }
                }
            }
        }
        
        /* Compress all chains. This is basically a union-find "path compression"
         * step.
         */
        function<string (string)> compress = [&](string label) {        
            if (label == newNames[label]) return label;
            return newNames[label] = compress(newNames[label]);
        };
        
        for (auto [ label, _ ]: newNames) {
            compress(label);
        }
        
        /* Make sure not to compact away "Start." That always has to be there. */
        newNames["Start"] = "Start";
        
        /* Now translate the labels. */
        stringstream builder;
        for (size_t i = 0; i < program.numLines(); i++) {
            builder << StatementToString().calculate(*GotoTranslator(newNames).calculate(*program.statement(i))) << '\n';
        }
        
        return Turing::Program(builder);
    }
    
    /* Every symbol our univeral TM is equipped to handle. */
    set<char> allLegalSymbols() {
        static set<char> result = [] {
            set<char> result;
            for (int i = 0; i < 128; i++) {
                if (i == 0 || i == '\n' || i == ' ' || i == '\t' || isprint(i)) {
                    result.insert(i);
                }
            }
            return result;
        } ();
        return result;
    }
    
    /* Calculator that produces the transfer function result for the constant
     * propagation optimization. The output from the calculation is map telling
     * us which lines will be reachable with which combinations of symbols.
     */
    class TransferFunctionCalculator: public Turing::Calculator<map<size_t, set<char>>> {
    public:
        TransferFunctionCalculator(const set<char>& inSet, size_t lineNumber, const Turing::Program& program):
            inSet(inSet), lineNumber(lineNumber), program(program) {
            // Handled in initializers
        }
        
        /* Labels push the information on entry to the next line. */
        map<size_t, set<char>> handle(Turing::Label&) override {
            return sendToNextLine(inSet);
        }
        
        /* Moves reset everything on the next line. */
        map<size_t, set<char>> handle(Turing::Move&) override {
            return sendToNextLine(allLegalSymbols());
        }
        
        /* Returns don't send anything anywhere. */
        map<size_t, set<char>> handle(Turing::Return&) override {
            return { };
        }
        
        /* Writes send the written character to the next line. */
        map<size_t, set<char>> handle(Turing::Write& w) override {            
            return sendToNextLine({ char(w.ch()) });
        }
        
        /* Gotos send the in set to the label target. */
        map<size_t, set<char>> handle(Turing::Goto& g) override {
            return { { program.lineForLabel(g.label()), inSet } };
        }
        
        /* If statements are a bit complex.
         *
         * The nested statement can either be a fallthrough statement
         * or a non-fallthrough statement (Goto or Return). If we have
         * a fallthrough statement, then we always proceed to the next
         * line, but in doing so need to consider the effects of both
         * executing and not executing the If statement. We thus union
         * together our in set and the result of our nested statement.
         * Importantly, our nested statement should only produce a
         * map of size 0 or 1 (0 if there is no next statement, 1
         * if there is) because we can't have any branching.
         *
         * If we have a non-fallthrough statement (Return or Goto),
         * then things get more interesting. Anything that makes it
         * to the next line of the program will be filtered down
         * based on the conditional. So if we have an If Not and we
         * don't take it, we can remove everything except for the
         * symbol in the conditional from the next line. If we have
         * an If and we don't take it, we remove that character when
         * pushing our in set down one line.
         *
         * Separately, we need to filter the effect of the Goto target.
         * Taking an If ... Goto means that Goto target is limited to
         * just the character in question (if it's even possible).
         * Tking an If Not ... Goto removes that character from the
         * in set.
         */
        map<size_t, set<char>> handle(Turing::If& i, const map<size_t, set<char>>& nestedResult) override {        
            /* We might be able to conclude that the If statement can't fire given the
             * information we have. If so, we proceed to the next line with nothing
             * changed.
             *
             * An If where the character in question isn't in the In set can't fire.
             */
            if (!i.isNegated() && !inSet.contains(i.ch())) {
                return sendToNextLine(inSet);
            }
            
            /* An If Not where there are no elements in the in set besides the character
             * in question can't fire.
             */
            if (i.isNegated() && inSet == set<char>{ char(i.ch()) }) {
                return sendToNextLine(inSet);
            }
        
            /* We now know that the If statement could in principle fire. Now we have to see what
             * happens if it does vs it doesn't.
             */
        
            /* First see what happens to the information that falls through to the next line. */
            set<char> toNextLine = inSet;
            
            /* Fallthrough statement. */
            if (typeid(*i.stmt()) != typeid(Turing::Goto) && 
                typeid(*i.stmt()) != typeid(Turing::Return)) {
                /* Could be empty if there's no next line. */
                if (!nestedResult.empty()) {
                    /* Should be exactly one thing, which is for the next line. */
                    if (nestedResult.size() != 1) throw runtime_error("Non-branching statement has multiple targets?");
                    
                    auto itr = nestedResult.begin();
                    if (itr->first != lineNumber + 1) throw runtime_error("Non-branching statement goes somewhere other than the next line?");
                    
                    /* Everything in here needs to be unioned in with what gets sent to the next line. */
                    toNextLine.insert(itr->second.begin(), itr->second.end());
                }
                
                return sendToNextLine(toNextLine);
            }
            /* Non-fallthrough statement. Update what goes to the next line based on what kind of
             * if statement this is.
             */
            else {
                /* If statement: we fell through because the If didn't trigger, so we
                 * can't have the given character.
                 */
                if (!i.isNegated()) {
                    toNextLine.erase(i.ch());
                }
                else {
                    /* If Not statement: we fell through because the If Not didn't trigger,
                     * so we must only have the given character (or nothing at all if
                     * that character isn't possible here.)
                     */
                    if (toNextLine.contains(i.ch())) {
                        toNextLine = { char(i.ch()) };
                    } else {
                        toNextLine = { };
                    }
                }
                
                /* We now know what's going to the next line. But now we need to
                 * trim down what's going to the branch target - if there is one.
                 */
                auto newDestinationSet = nestedResult;
                if (newDestinationSet.empty()) {
                    return sendToNextLine(toNextLine);
                }
                if (newDestinationSet.size() != 1) throw runtime_error("Nested branching statement has multiple targets?");
                
                /* This iterator points at the key/value pair of the destination line
                 * coupled with the characters that reach it.
                 */
                auto itr = newDestinationSet.begin();
                
                /* If we're an If Not statement that fires, the destination
                 * additions need to have the character in question removed.
                 */
                if (i.isNegated()) {
                    itr->second.erase(i.ch());
                }
                /* Otherwise, we're an If that fires. The destination set must
                 * become the character in question. (We handled the "this can't
                 * fire" case above.)
                 */
                else {
                    itr->second = { char(i.ch()) };
                }
                
                map<size_t, set<char>> result = sendToNextLine(toNextLine);
                result[itr->first].insert(itr->second.begin(), itr->second.end());
                return result;
            }
        }
    
    private:
        /* A map that sends the given data to the next line, if that next line
         * even exists.
         */
        map<size_t, set<char>> sendToNextLine(const set<char>& result) {
            /* If we hae a next line, send it the In set. */
            if (auto next = nextLine()) {
                return { { next.value(), result } };
            }
            /* Otherwise send nothing. */
            else {
                return { };
            }
        }
    
        /* Finds the next line after the given one. */
        optional<size_t> nextLine() const {
            if (lineNumber + 1 < program.numLines()) return lineNumber + 1;
            return nullopt;
        }
    
        set<char> inSet;   // What's alive on entry here?
        size_t lineNumber; // Which line are we on?
        const Turing::Program& program; // Which program are we in?
    };
    
    /* Computes the IN sets for each statement in the constant propagation
     * analysis. This yields, for each statement, a conservative overapproximation
     * of all symbols that could be under the tape head at the point that
     * the statement is executed.
     */
    vector<set<char>> inSetsFor(const Turing::Program& program) {
        /*
         * Initially everything is set to the bottom element (empty set),
         * except the Start label, which could begin with anything on
         * the tape.
         */
        vector<set<char>> inSets(program.numLines());
        inSets.at(program.lineForLabel("Start")) = allLegalSymbols();
        
        /* Keep running the optimization until it converges.
         *
         * TODO: The "right" way to do this is with a DFS postorder traversal,
         * which minimizes the number of passes needed. But any order works.
         */
        vector<set<char>> lastSets;
        do {
            /* Trach what we used to have to see if we've converged. */
            lastSets = inSets;
            
            for (size_t i = 0; i < program.numLines(); i++) {
                /* For each statement, see the effect of executing that
                 * statement. This will result in some number of symbols
                 * being propagated downward to the successors. The challenge
                 * is that different data might go to each successor. For
                 * example, the statement If 'a' Goto Label will send 'a' to
                 * Label but send Sigma - a to the next line. Thus each statement
                 * will produce a map<size_t, set<size_t>> showing what's sent
                 * to each statement. We then integrate them.
                 *
                 * If the In set here is empty, though, then the transfer function
                 * implicity returns the empty set and we don't need to do
                 * anything.
                 */
                if (inSets[i].empty()) continue;
                 
                /* Nonempty, so let's see what happens. */
                auto results = TransferFunctionCalculator(inSets[i], i, program).calculate(*program.statement(i));
                for (auto& [ line, set ]: results) {
                    inSets[line].insert(set.begin(), set.end());
                }
            }
        } while (lastSets != inSets);
        
        return inSets;
    }
    
    /* Given the In set for an If statement, sees if that If statement always
     * fires.
     */
    bool alwaysFires(Turing::If& i, const set<char>& inSet) {
        /* If Not always fires if the character isn't there. */
        if (i.isNegated()) {
            return !inSet.contains(i.ch());
        }
        /* If always fires if the only option is that character. */
        else {
            return inSet.size() == 1 && inSet.contains(i.ch());
        }
    }
    
    /* Given the In set for an If statement, sees if that If statement
     * never fires.
     */
    bool neverFires(Turing::If& i, const set<char>& inSet) {
        /* If Not never fires if the only option is the given character. */
        if (i.isNegated()) {
            return inSet.size() == 1 && inSet.contains(i.ch());
        }
        /* If never fires if the character isn't a possibility. */
        else {
            return !inSet.contains(i.ch());
        }
    }
    
    /* A twist on the classic constant propagation optimization used by many
     * compilers (see Stanford CS243 or the Purple Dragon Book). The twist
     * here is that we have no variables to work with. Instead, we'll do
     * propagation based on what's directly under the tape head. We can infer
     * what it is based on what we write and the results of If statements.
     *
     * We can use this to eliminate conditionals that can never fire,
     * remove the If from conditionals that always fire, and identify
     * unreachable code to strip out.
     */
    Turing::Program propagateConstants(const Turing::Program& program) {
        /* Our semilattice consists of sets of symbols that can be under
         * the tape head, ordered by inclusion.
         *
         * Compute the final IN sets for each line of code.
         */
        auto inSets = inSetsFor(program);

        /* Now that we have this information, we can perform various optimizations:
         *
         * 1. If [Not] Statements that can't fire can be removed.
         * 2. If [Not] Statements that always fire can be replaced by their payload.
         * 3. Anything unreachable can be deleted.
         */
        stringstream builder;
        for (size_t i = 0; i < program.numLines(); i++) {
            /* Skip unreachable code. */
            if (inSets[i].empty()) continue;
            
            /* If statements are special. */
            if (auto ifS = dynamic_pointer_cast<Turing::If>(program.statement(i))) {
                if (alwaysFires(*ifS, inSets[i])) {
                    builder << StatementToString().calculate(*ifS->stmt()) << '\n';
                } else if (!neverFires(*ifS, inSets[i])) {
                    builder << program.line(i) << '\n';
                }
            }
            /* Anything else gets output as before. */
            else {
                builder << program.line(i) << '\n';
            }
        }      
        return Turing::Program(builder);
    }
    
    /* Removes all blank lines from the given program. */
    Turing::Program trimBlankLinesFrom(const Turing::Program& program) {
        stringstream result;
        for (size_t i = 0; i < program.numLines(); i++) {
            /* See if there's a statement here; that's a much better way than
             * string processing to see if the line is blank.
             */
            if (program.statement(i)) {
                result << program.line(i) << endl;
            }
        }
        
        return Turing::Program(result);
    }
}

int main(int argc, const char* argv[]) try {
    auto args = parseArguments(argc, argv);

    Turing::Program program(cin);
    
    /* Display errors if program can't be run. */
    if (!program.isValid()) {
        showErrorsIn(program);
        return -1;
    }
    
    /* Remove all blank lines from the program; this allows us to make the
     * assumption that line == statement.
     */
    program = trimBlankLinesFrom(program);
    
    /* Remember how many lines we started with. */
    size_t numLines;
    
    /* Run optimization in a loop as the number of lines decreases. */
    do {
        numLines = program.numLines();
        
        /* Some labels consist purely of a single unconditional Goto.
         * Rewrite the program to have code that jumps to those labels
         * instead jump to the ultimate target. (We could do some
         * more aggressive work here, but since we're already going to
         * remove unused labels in the next step, that's not necessary.)
         */
        program = bypassUnnecessaryLabels(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Bypass unnecessary labels." << endl;
            showErrorsIn(program);
            exit(-1);
        }
    
        /* Remove all labels that aren't used as Goto targets.
         * This also will remove all whitespace lines and comments,
         * so we're now left with something where line == statement.
         */
        program = removeUnusedLabelsFrom(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Remove unused labels." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Rewrite the program, reordering basic blocks to eliminate unneeded Gotos and labels. */
        program = reorderCode(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Reorder code." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Keyhole optimization: remove any Gotos that go immediately to the next line. */
        program = removeUnneededGotos(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Remove unneeded gotos." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Consecutive labels are never needed. */
        program = removeConsecutiveLabels(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Remove consecutive labels." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Look for consecutive moves going opposite directions. */
        program = removeUselessMoves(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Remove useless moves." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Apply constant propagation based on our knowledge of
         * what's immediately under the tape head.
         */
        program = propagateConstants(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Constant propagation." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
        /* Keyhole optimization: Collapse chains of Writes into
         * just the last Write.
         */
        program = removeConsecutiveWrites(program);
        if (!program.isValid()) {
            cerr << "OPTIMIZATION ERROR: Remove consecutive writes." << endl;
            showErrorsIn(program);
            exit(-1);
        }
        
    } while (program.numLines() < numLines);
    
    if (args.simplifyNames) {
        program = simplifyNames(program);
    }
    
    /* Output the resulting program. */
    for (size_t i = 0; i < program.numLines(); i++) {
        cout << program.line(i) << '\n';
    }  
    
} catch (const exception& e) {
    cerr << "Error: " << e.what() << endl;
    return -1;
}