reuse block key work in transport hot path

Author: OpusNano

src/core/energy.zig

@@ -5,11 +5,20 @@ pub fn pairWeight(comptime T: type, left: T, right: T, cfg: Config) u64 {
     return 1 + key.closenessBonus(T, left, right, cfg);
 }
 
+pub fn pairWeightFromKeys(comptime T: type, left_key: key.KeyType(T), right_key: key.KeyType(T), cfg: Config) u64 {
+    return 1 + key.closenessBonusFromKeys(T, left_key, right_key, cfg);
+}
+
 pub fn pairEnergy(comptime T: type, left: T, right: T, cfg: Config) u64 {
     if (!key.greaterThan(T, left, right)) return 0;
     return pairWeight(T, left, right, cfg);
 }
 
+pub fn pairEnergyFromKeys(comptime T: type, left_key: key.KeyType(T), right_key: key.KeyType(T), cfg: Config) u64 {
+    if (!key.greaterThanKeys(T, left_key, right_key)) return 0;
+    return pairWeightFromKeys(T, left_key, right_key, cfg);
+}
+
 pub fn blockEnergy(comptime T: type, xs: []const T, cfg: Config) u64 {
     var total: u64 = 0;
     for (xs, 0..) |left, i| {
@@ -20,3 +29,14 @@ pub fn blockEnergy(comptime T: type, xs: []const T, cfg: Config) u64 {
     }
     return total;
 }
+
+pub fn blockEnergyFromKeys(comptime T: type, keys: []const key.KeyType(T), cfg: Config) u64 {
+    var total: u64 = 0;
+    for (keys, 0..) |left_key, i| {
+        var j = i + 1;
+        while (j < keys.len) : (j += 1) {
+            total += pairEnergyFromKeys(T, left_key, keys[j], cfg);
+        }
+    }
+    return total;
+}

src/core/key.zig

@@ -2,6 +2,8 @@ const std = @import("std");
 const Config = @import("config.zig").Config;
 const types = @import("types.zig");
 
+pub const KeyType = types.UnsignedOf;
+
 pub fn biasedKey(comptime T: type, value: T) types.UnsignedOf(T) {
     const info = @typeInfo(T);
     const U = types.UnsignedOf(T);
@@ -20,8 +22,12 @@ pub fn biasedKey(comptime T: type, value: T) types.UnsignedOf(T) {
 }
 
 pub fn closenessBonus(comptime T: type, left: T, right: T, cfg: Config) u8 {
+    return closenessBonusFromKeys(T, biasedKey(T, left), biasedKey(T, right), cfg);
+}
+
+pub fn closenessBonusFromKeys(comptime T: type, left_key: types.UnsignedOf(T), right_key: types.UnsignedOf(T), cfg: Config) u8 {
     const U = types.UnsignedOf(T);
-    const diff: U = biasedKey(T, left) ^ biasedKey(T, right);
+    const diff: U = left_key ^ right_key;
     if (diff == 0) return cfg.valuation_cap;
 
     const tz: usize = @ctz(diff);
@@ -37,3 +43,7 @@ pub fn compare(comptime T: type, left: T, right: T) std.math.Order {
 pub fn greaterThan(comptime T: type, left: T, right: T) bool {
     return left > right;
 }
+
+pub fn greaterThanKeys(comptime T: type, left_key: types.UnsignedOf(T), right_key: types.UnsignedOf(T)) bool {
+    return left_key > right_key;
+}

src/core/pressure.zig

@@ -30,6 +30,29 @@ pub fn compute(comptime T: type, block: []const T, cfg: Config, pressures: []i32
     }
 }
 
+pub fn computeFromKeys(comptime T: type, keys: []const key.KeyType(T), cfg: Config, pressures: []i32) void {
+    std.debug.assert(pressures.len >= keys.len);
+
+    @memset(pressures[0..keys.len], 0);
+    if (keys.len <= 1) return;
+
+    for (keys, 0..) |left_key, i| {
+        const end = @min(keys.len, i + cfg.neighborhood + 1);
+        var j = i + 1;
+        while (j < end) : (j += 1) {
+            const right_key = keys[j];
+            if (left_key > right_key) {
+                const w: i32 = @intCast(energy.pairWeightFromKeys(T, left_key, right_key, cfg));
+                pressures[i] += w;
+                pressures[j] -= w;
+            } else if (left_key < right_key) {
+                pressures[i] -= 1;
+                pressures[j] += 1;
+            }
+        }
+    }
+}
+
 pub fn proposalsFromPressure(pressures: []const i32, cfg: Config, proposals: []i8) void {
     std.debug.assert(proposals.len >= pressures.len);
 

src/core/transport.zig

@@ -1,8 +1,10 @@
 const std = @import("std");
 const Config = @import("config.zig").Config;
 const energy = @import("energy.zig");
+const key = @import("key.zig");
 const pressure = @import("pressure.zig");
 const Stats = @import("stats.zig").Stats;
+const util = @import("../internal/util.zig");
 
 pub const BlockResult = struct {
     accepted: bool,
@@ -58,8 +60,16 @@ pub fn tryTransportBlock(comptime T: type, block: []T, cfg: Config, stats: ?*Sta
     std.debug.assert(block.len <= Config.max_block_size);
     if (stats) |s| s.transport_blocks_visited += 1;
 
-    const before_energy = energy.blockEnergy(T, block, cfg);
-    if (block.len <= 1 or before_energy == 0) {
+    if (block.len <= 1 or util.isSorted(T, block)) {
+        if (stats) |s| s.transport_blocks_rejected += 1;
+        return .{ .accepted = false, .before_energy = 0, .after_energy = 0 };
+    }
+
+    var keys: [Config.max_block_size]key.KeyType(T) = undefined;
+    for (block, 0..) |value, i| keys[i] = key.biasedKey(T, value);
+
+    const before_energy = energy.blockEnergyFromKeys(T, keys[0..block.len], cfg);
+    if (before_energy == 0) {
         if (stats) |s| s.transport_blocks_rejected += 1;
         return .{ .accepted = false, .before_energy = before_energy, .after_energy = before_energy };
     }
@@ -72,8 +82,9 @@ pub fn tryTransportBlock(comptime T: type, block: []T, cfg: Config, stats: ?*Sta
     var desired: [Config.max_block_size]usize = undefined;
     var source_to_final: [Config.max_block_size]usize = undefined;
     var candidate: [Config.max_block_size]T = undefined;
+    var candidate_keys: [Config.max_block_size]key.KeyType(T) = undefined;
 
-    pressure.compute(T, block, cfg, pressures[0..block.len]);
+    pressure.computeFromKeys(T, keys[0..block.len], cfg, pressures[0..block.len]);
     pressure.proposalsFromPressure(pressures[0..block.len], cfg, proposals[0..block.len]);
 
     var changed = false;
@@ -95,9 +106,10 @@ pub fn tryTransportBlock(comptime T: type, block: []T, cfg: Config, stats: ?*Sta
     std.debug.assert(isPermutation(source_to_final[0..block.len]));
     for (block, 0..) |value, source_index| {
         candidate[source_to_final[source_index]] = value;
+        candidate_keys[source_to_final[source_index]] = keys[source_index];
     }
 
-    const after_energy = energy.blockEnergy(T, candidate[0..block.len], cfg);
+    const after_energy = energy.blockEnergyFromKeys(T, candidate_keys[0..block.len], cfg);
     if (after_energy >= before_energy) {
         std.debug.assert(equalSlices(T, block, original[0..block.len]));
         if (stats) |s| s.transport_blocks_rejected += 1;