Heuristically choose and preserve rdx across mulx

src/assembler/assembler.helper.ts

@@ -15,7 +15,7 @@
  */
 
 import { Flags, Register, XmmRegister } from "@/enums";
-import { getQwRegFromByteReg, isByteRegister, isFlag, isRegister, isXmmRegister, matchArg } from "@/helper";
+import { getQwRegFromByteReg, isByteRegister, isFlag, isRegister, isMem, isXmmRegister, matchArg } from "@/helper";
 import { RegisterAllocator } from "@/registerAllocator";
 import type { CryptOpt } from "@/types";
 
@@ -51,7 +51,7 @@ export function sanityCheckAllocations(c: CryptOpt.DynArgument): void {
         }
         byReg[r64] = varname;
       }
-      if (matchArg(varname)) {
+      if (matchArg(varname) && isMem(store)) {
         throw new Error("should not be allocated.");
       }
       return byReg;

src/instructionGeneration/multiplication.ts

@@ -148,9 +148,11 @@ function mulx64(ra: RegisterAllocator, c: CryptOpt.StringOperation): asm[] {
   });
 
   const [resLoR, resHiR] = allocation.oReg;
-  const [arg0R, arg1R] = allocation.in;
-  let argR = arg0R !== Register.rdx ? arg0R : arg1R;
-
+  const explicitArgs = allocation.in.filter(r => r !== Register.rdx);
+  if (explicitArgs.length !== 1) {
+    throw new Error("TSNH. mulx can only take one argument besides rdx");
+  }
+  let [argR] = explicitArgs;
   argR = makeArgRanR64(argR, ra);
 
   // if can use mulx

src/model/model.class.ts

@@ -25,6 +25,7 @@ import {
   DI_ABBRV,
   isADependentOnB,
   isCallerSave,
+  isXD,
   limbify,
   matchArg,
   matchArgPrefix,
@@ -203,9 +204,54 @@ export class Model {
   }
 
   /*
-   * @param candidates is a list of variable names like xNN, argNN, ...
-   * returns the one that we, in the current ordering, read last.
+   * @param candidates is a list of variable names like xNN, argNN, that are arguments to the curernt mulx
+   * returns the one that is loaded most in upcoming mulxs, and also returns a bit of debug inforamtion
    */
+  public static chooseMulxLoadValue(candidates: string[]): { msg: string; candidate: string | null } {
+    let msg = "; chooseMulxLoadValue";
+    if (candidates.length < 1) {
+      throw new Error("cannot choose from nothing, mate");
+    }
+    // well, not much to choose from, right?
+    if (candidates.length === 1) {
+      const [candidate] = candidates;
+      return { candidate, msg };
+    }
+
+    // Idea is:
+    // for each candidate, count how many consecutive upcoming mulx operations use it
+
+    // initalise counters:
+    const counters: { [candidateName: string]: number } = candidates.reduce(
+      (acc, candidate) => Object.assign(acc, { [candidate]: 0 }),
+      {},
+    );
+    const upcomingMulxOps = Model.nodesInTopologicalOrder
+      .slice(this._currentInstIdx)
+      .filter((op) => op.operation == "mulx");
+    msg += ` upcoming mulxs: ${upcomingMulxOps.map((m) => m.name.join("-")).join(", ")}`;
+
+    candidates.forEach((arg) => {
+      upcomingMulxOps.every((op) => {
+	if (!op.arguments.includes(arg)) { return false; }
+	counters[arg]++;
+	return true;
+      });
+    });
+    msg += ` counters ${JSON.stringify(counters)}`;
+
+    const sortedCounters = Object.entries(counters)
+      .sort(([, a], [, b]) => a - b)
+      .reverse();
+    // and return it's name
+    // if there is only two, and they are the same
+    if (sortedCounters.length == 2 && sortedCounters[0][1] == sortedCounters[1][1]) {
+      msg += ` only two candidates, and they both have the same count value of ${sortedCounters[0][1]}. returning null.`;
+      return { candidate: null, msg };
+    }
+    msg += ` choosing ${sortedCounters[0][0]} because of its higher count value of ${sortedCounters[0][1]}`;
+    return { candidate: sortedCounters[0][0], msg };
+  }
   public static chooseSpillValue(candidates: string[]): string {
     if (candidates.length < 1) {
       throw new Error("cannot choose from nothing, mate");

src/paul/Paul.class.ts

@@ -175,6 +175,8 @@ export class Paul {
       return Math.floor(stateBasedFactor * delta) + start;
     }
 
+    // https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
+    // and https://stackoverflow.com/questions/25582882/javascript-math-random-normal-distribution-gaussian-bell-curve
     const randn_bm = (min: number, max: number, skew: number): number => {
       let u = 0,
         v = 0;

src/registerAllocator/RegisterAllocator.class.ts

@@ -463,6 +463,9 @@ export class RegisterAllocator {
     const caf = (flagToCheck: AllocationFlags): boolean =>
       ((allocationReq.allocationFlags ?? AllocationFlags.NONE) & flagToCheck) === flagToCheck;
     const inAllocationsTemp = allocationReq.in.map((readVariable) => {
+      const currentLocation = this._allocations[readVariable];
+      if (currentLocation && isRegister(currentLocation.store)) return currentLocation.store;
+
       const argMatchRes = matchArg(readVariable);
       if (argMatchRes) {
         // if we read from an argument such as arg1[3], we actually want to find arg1 in the allocations.
@@ -474,7 +477,6 @@ export class RegisterAllocator {
           throw new Error(`${readVariable} matched ~arg, but has no baseVar? wtf. Giving up.`);
         }
       }
-      const currentLocation = this._allocations[readVariable];
       if (!currentLocation) {
         // if is it not already allocated, it must be an immval, cause 'arg1' is always somewhere
         if (caf(AllocationFlags.DISALLOW_IMM)) {
@@ -593,7 +595,7 @@ export class RegisterAllocator {
       }
     }
 
-    if (caf(AllocationFlags.ONE_IN_MUST_BE_IN_RDX) && !inAllocations.includes(Register.rdx)) {
+    if (caf(AllocationFlags.ONE_IN_MUST_BE_IN_RDX) && !allocationReq.in.some((i) => this._allocations[i].store === Register.rdx)) {
       // since in inAllocations there is no rdx (otherwise we wont be in this branch)
       // we need to move one of inAllocations to rdx
 
@@ -605,8 +607,17 @@ export class RegisterAllocator {
 
       // we want now change any of those inAllocations with rdx.
 
-      // Paul chooses an element, which we'll move to rdx.
-      const element = Paul.chooseArg(allocationReq.in);
+      // try to be sophisticated in choosing the mulx Load value
+      const { msg, candidate } = Model.chooseMulxLoadValue(allocationReq.in);
+      this.addToPreInstructions(msg);
+      let element = "";
+      if (candidate != null) {
+        element = candidate;
+      } else {
+        // if this didn't work (none is clearly preferrable)
+        // let Paul choose an element, which we'll move to rdx.
+        element = Paul.chooseArg(allocationReq.in);
+      }
       const idx = allocationReq.in.indexOf(element);
       //TODO: refactor that a lil bit
 
@@ -1500,7 +1511,9 @@ export class RegisterAllocator {
     Object.keys(this._allocations)
       .filter(matchArg)
       .forEach((v) => {
-        delete this._allocations[v];
+        if (isMem(this._allocations[v].store)) {
+          delete this._allocations[v];
+        }
       });
 
     // empty and get current pres