day 17 python refactor

day17/sol.py

@@ -16,39 +16,46 @@ a = 0b11111111
 wall = 0b10000000_10000000_10000000_10000000_10000000
 stopped = 0
 r = 0
-while stopped < 2022:
-    if r == 0:
-        # start falling
-        r = next(rocki)
-        h = a.bit_length() + 8*3
-        if h % 8 != 0:
-            h += 8 - h%8
-        #for i in reversed(range(0,a.bit_length(),8)):
-        #    print("{:08b}".format((a>>i)&0xff))
-        #print()
-    else:
+
+def drop(a, r, jets):
+    # start falling
+    h = a.bit_length() + 8*3
+    if h % 8 != 0:
+        h += 8 - h%8
+    while True:
+        j = next(jets)
+        if j == '<':
+            x = (r<<1)
+            if x & wall == 0:
+                if x & (a>>h) == 0:
+                    r = x
+        elif j == '>':
+            x = (r<<7)
+            if x & wall == 0:
+                x >>= 8
+                if x & (a>>h) == 0:
+                    r = x
+
         # continue falling?
         if r & (a>>(h-8)) == 0:
             h -= 8
+            assert h > 0
         else:
             # stop
             a |= (r<<h)
-            r = 0
-            stopped += 1
-            continue
+            return a
 
-    j = next(jets)
-    if j == '<':
-        x = (r<<1)
-        if x & wall == 0:
-            if x & (a>>h) == 0:
-                r = x
-    elif j == '>':
-        x = (r<<7)
-        if x & wall == 0:
-            x >>= 8
-            if x & (a>>h) == 0:
-                r = x
+
+a = 0b11111111
+wall = 0b10000000_10000000_10000000_10000000_10000000
+stopped = 0
+while stopped < 2022:
+    r = next(rocki)
+    a = drop(a,r,jets)
+    stopped += 1
+    #for i in reversed(range(0,a.bit_length(),8)):
+    #    print("{:08b}".format((a>>i)&0xff))
+    #print()
 
 import math
 print(math.ceil((a>>8).bit_length()/8))