Solution to Advent of Code 2020 Day 14 Part 2

AOC_2020_Day14_Part2.js

const { parseInputFile } = require('../utils/parser');

/** @typedef {'0' | '1' | 'X'} Trit */

/**
 * @param {string} addr 
 */
function parseAddr(addr) {
    const arr = Array.from((parseInt(addr) >>> 0).toString(2));
    while (arr.length < 36) {
        arr.unshift('0');
    }
    return arr;
}

/** @type {Array<{ ins: string, addr?: Trit[], mask?: Trit[], value?: bigint}>} */
// @ts-ignore
const inputs = parseInputFile('day14-input.txt', /^(mask|mem)(?:\[(\d+)\])? = ([0-9X]+)$/gm)
    .map(([, ins, addr, maskOrValue]) => {
        switch (ins) {
            case 'mask': return {
                ins,
                mask: Array.from(maskOrValue)
            };
            case 'mem': return {
                ins,
                addr: parseAddr(addr),
                value: BigInt(parseInt(maskOrValue))
            };
            default: throw new Error(`Invalid instruction: ${ ins }`);
        }
    });

/**
 * @typedef {object} MemoryNode
 * @property {Trit} addr
 * @property {boolean} isLeaf
 * @property {MemoryNode[]} nextNodes
 * @property {bigint} [value]
 * @property {boolean} isRoot
 */

/** @type {MemoryNode} */
const memoryRoot = {
    addr: 'X',
    isLeaf: false,
    nextNodes: [],
    isRoot: true
};

/**
 * @param {MemoryNode} node 
 * @param {Trit[]} addr 
 * @param {number} nextAddrOffset
 * @param {bigint} value
 * @param {Trit} nextAddr
 */
function handleBranchNode(node, addr, nextAddrOffset, value, nextAddr) {
    // find the node (if it exists) that matches the address
    const matchingNode = node.nextNodes.find(nn => nn.addr === nextAddr);
    if (matchingNode !== undefined) {
        // Option 2: node NOT leaf, trit match
        writeValueAt(matchingNode, addr, nextAddrOffset, value);

    } else {
        // Option 3: node NOT leaf, trit NOT match

        // address of next node
        const nextNodeIsLeaf = addr.length - nextAddrOffset <= 1;

        // create next node
        const nextNode = {
            addr: nextAddr,
            isLeaf: nextNodeIsLeaf,
            nextNodes: [],
            isRoot: false
        };

        // attach and recurse
        node.nextNodes.push(nextNode);
        writeValueAt(nextNode, addr, nextAddrOffset, value);
    }
}

/**
 * @param {MemoryNode} node 
 * @param {Trit[]} addr 
 * @param {number} addrOffset
 * @param {bigint} value 
 */
function writeValueAt(node, addr, addrOffset, value) {
    // zero needs to be repeated, since there is N + 1 nodes
    const nextAddrOffset = node.isRoot ? addrOffset : addrOffset + 1;

    // Option 1: node IS leaf
    if (node.isLeaf) {
        // write value
        node.value = value;

    } else {
        const nextAddr = addr[nextAddrOffset];
        switch (nextAddr) {
            case '0': 
            case '1':
                handleBranchNode(node, addr, nextAddrOffset, value, nextAddr);
                break;
            case 'X':
                handleBranchNode(node, addr, nextAddrOffset, value, '0');
                handleBranchNode(node, addr, nextAddrOffset, value, '1');
                break;
        }
    }
}

/**
 * @param {Trit[]} address 
 * @param {bigint} value 
 */
function writeValue(address, value) {
    writeValueAt(memoryRoot, address, 0, value);
}

/**
 * @param {MemoryNode} node 
 * @param {string} addr 
 * @param {Array<{ addr: string, value: bigint }>} values 
 */
function addMemoryValues(node, addr, values) {
    if (!node.isRoot) {
        addr += node.addr;
    }

    if (node.isLeaf) {
        values.push({
            addr: addr,
            value: node.value
        });
    } else {
        for (const nextNode of node.nextNodes) {
            addMemoryValues(nextNode, addr, values);
        }
    }
}
/**
 * @returns {Array<{ addr: string, value: bigint }}
 */
function getMemoryValues() {
    const values = [];
    addMemoryValues(memoryRoot, '', values);
    return values;
}

// pre-initialize empty mask
/** @type {Trit[]} */
let mask = ['0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0'];

// emulate all instructions
for (const instruction of inputs) {
    if (instruction.ins === 'mask') {
        // update mask
        mask = instruction.mask;

    } else {
        // set a value
        const addr = instruction.addr.map((bit, i) => (mask[i] === '0') ? bit : mask[i]);
        writeValue(addr, instruction.value);
    }
}

const part2Answer = getMemoryValues().reduce((tot, entry) => tot + entry.value, 0n);
console.log(`Part 2: ${ part2Answer }`);