Day 22: adventures in modular math

1 files changed, 148 insertions, 86 deletions

diff --git a/src/bin/day_22.rs b/src/bin/day_22.rs
index fbe3f4c..5b999a6 100644
--- a/src/bin/day_22.rs
+++ b/src/bin/day_22.rs
@@ -1,6 +1,8 @@
+use derive_more::Display;
 use num::bigint::BigInt;
 use num::traits::identities::Zero;
 use num::traits::sign::abs;
+use num::traits::Signed;
 use std::fmt;
 use std::io;
 use std::io::prelude::*;
@@ -19,7 +21,7 @@ struct Opt {
     /// At the end, query the position of card
     card: BigInt,
     /// Number of repetitions
-    repetitions: usize,
+    repetitions: BigInt,
 
     /// Prints the card in position n, rather than the position of card n
     #[structopt(short = "p")]
@@ -42,22 +44,30 @@ fn main() {
     if opt.position_mode {
         println!(
             "{}",
-            Card {
-                position: opt.card,
-                modulus: opt.deck_size
-            }
-            .shuffle_reverse(&instructions, opt.repetitions)
-            .position
+            instructions
+                .iter()
+                .rev()
+                .fold(
+                    StandardisedInstruction::identity(opt.deck_size.clone()),
+                    |acc, next| acc.then(&(next.clone(), opt.deck_size.clone(), false).into())
+                )
+                .repeat(opt.repetitions)
+                .apply(opt.card.clone())
         );
     } else {
         println!(
             "{}",
-            Card {
-                position: opt.card,
-                modulus: opt.deck_size
-            }
-            .shuffle(&instructions, opt.repetitions)
-            .position
+            instructions
+                .iter()
+                .fold(
+                    StandardisedInstruction::identity(opt.deck_size.clone()),
+                    |acc, next| {
+                        eprintln!("{}", acc);
+                        acc.then(&(next.clone(), opt.deck_size.clone(), true).into())
+                    }
+                )
+                .repeat(opt.repetitions)
+                .apply(opt.card.clone())
         );
     }
 }
@@ -72,93 +82,37 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
     }
 }
 
-#[derive(Clone)]
-struct Card {
-    position: BigInt,
-    modulus: BigInt,
-}
-
-impl Card {
-    fn shuffle(&self, instructions: &[Instruction], repetitions: usize) -> Card {
-        (0..repetitions).fold(self.clone(), |acc, _| {
-            instructions.iter().fold(acc, |card, instruction| {
-                //eprintln!("{} :: {}", deck.zero, deck.cards);
-                card.single_shuffle(instruction)
-            })
-        })
-    }
-
-    fn single_shuffle(&self, instruction: &Instruction) -> Card {
-        match instruction {
-            Instruction::DealIntoNewStack => Card {
-                position: self.modulus.clone() - self.position.clone() - 1,
-                modulus: self.modulus.clone(),
-            },
-            Instruction::Cut(n) => Card {
-                position: mod_sub(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-            Instruction::ReverseCut(n) => Card {
-                position: mod_plus(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-            Instruction::DealWithIncrement(n) => Card {
-                position: mod_times(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-        }
-    }
-
-    fn shuffle_reverse(&self, instructions: &[Instruction], repetitions: usize) -> Card {
-        (0..repetitions).fold(self.clone(), |acc, _| {
-            instructions.iter().rev().fold(acc, |card, instruction| {
-                //eprintln!("{} :: {}", deck.zero, deck.cards);
-                card.single_shuffle_reverse(instruction)
-            })
-        })
-    }
-
-    fn single_shuffle_reverse(&self, instruction: &Instruction) -> Card {
-        match instruction {
-            Instruction::DealIntoNewStack => Card {
-                position: self.modulus.clone() - self.position.clone() - 1,
-                modulus: self.modulus.clone(),
-            },
-            Instruction::Cut(n) => Card {
-                position: mod_plus(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-            Instruction::ReverseCut(n) => Card {
-                position: mod_sub(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-            Instruction::DealWithIncrement(n) => Card {
-                position: mod_divide(self.position.clone(), n.clone(), self.modulus.clone()),
-                modulus: self.modulus.clone(),
-            },
-        }
-    }
-}
-
 fn mod_plus(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
-    (a + b) % modulus
+    mod_normalize(a + b, modulus)
 }
 
 fn mod_sub(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
-    (a + modulus.clone() - b) % modulus
+    mod_normalize(a - b, modulus)
 }
 
 fn mod_times(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
-    (a * b) % modulus
+    mod_normalize(a * b, modulus)
 }
 
 fn mod_divide(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
     mod_times(a, mod_inverse(b, modulus.clone()), modulus)
 }
 
+fn mod_pow(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
+    a.modpow(&b, &modulus)
+}
+
+fn mod_normalize(a: BigInt, modulus: BigInt) -> BigInt {
+    if a.is_negative() {
+        a.clone() + modulus.clone() * (1 + abs(a) / modulus)
+    } else {
+        a % modulus
+    }
+}
+
 // NB: This may give nonsense if modulus isn't coprime with a
 fn mod_inverse(a: BigInt, modulus: BigInt) -> BigInt {
-    (euclid_gcd_coefficients(a, modulus.clone()).0 + modulus.clone()) as BigInt % modulus
+    mod_normalize(euclid_gcd_coefficients(a, modulus.clone()).0, modulus)
 }
 
 fn euclid_gcd_coefficients(a: BigInt, b: BigInt) -> (BigInt, BigInt) {
@@ -189,7 +143,7 @@ fn euclid_gcd_coefficients(a: BigInt, b: BigInt) -> (BigInt, BigInt) {
     euclid_gcd_coefficients_inner(b, a, 0.into(), 1.into(), 1.into(), 0.into())
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 enum Instruction {
     DealIntoNewStack,
     Cut(BigInt),
@@ -235,6 +189,114 @@ impl FromStr for Instruction {
     }
 }
 
+// f(x) = ax + b mod c
+#[derive(Display, Clone)]
+#[display(fmt = "f(x) = {} x + {} % {}", a, b, modulus)]
+struct StandardisedInstruction {
+    a: BigInt,
+    b: BigInt,
+    modulus: BigInt,
+}
+
+impl From<(Instruction, BigInt, bool)> for StandardisedInstruction {
+    fn from((instruction, modulus, forward): (Instruction, BigInt, bool)) -> Self {
+        match (instruction, forward) {
+            (Instruction::DealIntoNewStack, _) => StandardisedInstruction {
+                a: BigInt::from(-1),
+                b: BigInt::from(-1),
+                modulus: modulus,
+            },
+            (Instruction::Cut(n), true) => StandardisedInstruction {
+                a: BigInt::from(1),
+                b: BigInt::from(-n),
+                modulus: modulus,
+            },
+            (Instruction::Cut(n), false) => StandardisedInstruction {
+                a: BigInt::from(1),
+                b: BigInt::from(n),
+                modulus: modulus,
+            },
+            (Instruction::ReverseCut(n), true) => StandardisedInstruction {
+                a: BigInt::from(1),
+                b: BigInt::from(n),
+                modulus: modulus,
+            },
+            (Instruction::ReverseCut(n), false) => StandardisedInstruction {
+                a: BigInt::from(1),
+                b: BigInt::from(-n),
+                modulus: modulus,
+            },
+            (Instruction::DealWithIncrement(n), true) => StandardisedInstruction {
+                a: BigInt::from(n),
+                b: BigInt::from(0),
+                modulus: modulus,
+            },
+            (Instruction::DealWithIncrement(n), false) => StandardisedInstruction {
+                a: BigInt::from(mod_inverse(n, modulus.clone())),
+                b: BigInt::from(0),
+                modulus: modulus,
+            },
+        }
+        .normalise()
+    }
+}
+
+impl StandardisedInstruction {
+    fn identity(modulus: BigInt) -> StandardisedInstruction {
+        StandardisedInstruction {
+            a: BigInt::from(1),
+            b: BigInt::from(0),
+            modulus,
+        }
+    }
+    fn normalise(&self) -> StandardisedInstruction {
+        StandardisedInstruction {
+            a: mod_normalize(self.a.clone(), self.modulus.clone()),
+            b: mod_normalize(self.b.clone(), self.modulus.clone()),
+            modulus: self.modulus.clone(),
+        }
+    }
+    fn then(&self, other: &StandardisedInstruction) -> StandardisedInstruction {
+        // g(f(x)) = ga (fa x + fb) + gb =
+        StandardisedInstruction {
+            a: mod_times(self.a.clone(), other.a.clone(), self.modulus.clone()),
+            b: mod_plus(
+                mod_times(self.b.clone(), other.a.clone(), self.modulus.clone()),
+                other.b.clone(),
+                self.modulus.clone(),
+            ),
+            modulus: self.modulus.clone(),
+        }
+    }
+    fn repeat(&self, repetitions: BigInt) -> StandardisedInstruction {
+        StandardisedInstruction {
+            a: mod_pow(self.a.clone(), repetitions.clone(), self.modulus.clone()),
+            b: mod_divide(
+                mod_times(
+                    self.b.clone(),
+                    mod_sub(
+                        BigInt::from(1),
+                        mod_pow(self.a.clone(), repetitions.clone(), self.modulus.clone()),
+                        self.modulus.clone(),
+                    ),
+                    self.modulus.clone(),
+                ),
+                mod_sub(BigInt::from(1), self.a.clone(), self.modulus.clone()),
+                self.modulus.clone(),
+            ),
+            modulus: self.modulus.clone(),
+        }
+    }
+
+    fn apply(&self, x: BigInt) -> BigInt {
+        mod_plus(
+            mod_times(self.a.clone(), x, self.modulus.clone()),
+            self.b.clone(),
+            self.modulus.clone(),
+        )
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 struct ParseErr;