From ccfd03e401f1dd5ceddedae8c751328e1d7321ef Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 22 Dec 2020 10:16:46 -0500
Subject: [PATCH] parse ane struct

---
 ane/2_compile/hwx_parse.py      |  5 +++--
 ane/2_compile/struct_recover.py | 33 +++++++++++++++++++++++++++++++++
 ane/lib/ane.py                  | 13 ++++++++++---
 3 files changed, 46 insertions(+), 5 deletions(-)
 create mode 100755 ane/2_compile/struct_recover.py

diff --git a/ane/2_compile/hwx_parse.py b/ane/2_compile/hwx_parse.py
index 30288941d6..44768cb2f4 100755
--- a/ane/2_compile/hwx_parse.py
+++ b/ane/2_compile/hwx_parse.py
@@ -109,8 +109,9 @@ f1 = g.headers[0].commands[1][2][0].section_data
 f2 = a.headers[0].commands[1][2][0].section_data
 for i in range(0, len(f2), 0x300):
   print("===== op %d =====" % (i//0x300))
-  dbg = ane.debug(f2[i:i+0x300])
-  print(dbg)
+  dbg = ane.debug(f1[i:i+0x300], 2)
+  for k,v in dbg.items():
+    print(k, v)
   if len(f1) < 0x300:
     print(compare(f1, f2[i:i+0x300]))
   else:
diff --git a/ane/2_compile/struct_recover.py b/ane/2_compile/struct_recover.py
new file mode 100755
index 0000000000..e026ecb6dc
--- /dev/null
+++ b/ane/2_compile/struct_recover.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+from ane import ANE
+ane = ANE()
+
+lens = {}
+pos = {}
+
+dat = b"\xff"*0x300
+ret = ane.debug(dat, 16)
+for k,v in ret.items():
+  found = None
+  for i in range(33):
+    #print(v, (1 << i) - 1)
+    if v == (1 << i) - 1:
+      found = i
+      break
+  #print(k, hex(v), found)
+  lens[k] = found
+
+dat = b"\x00"*0x300
+for i in range(0x300):
+  for j in range(8):
+    dat = b"\x00"*i
+    dat += bytes([1 << j])
+    dat += b"\x00"*(0x300-len(dat))
+    ret = ane.debug(dat, 16)
+    for k,v in ret.items():
+      if v == 1:
+        print("0x%3x %d %2d" % (i, j, lens[k]), k)
+        pos[k] = (i,j, lens[k])
+
+print(len(lens), len(pos))
+
diff --git a/ane/lib/ane.py b/ane/lib/ane.py
index 5453628592..2498b48f2f 100755
--- a/ane/lib/ane.py
+++ b/ane/lib/ane.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import os
 from ctypes import *
+import collections
 import numpy as np
 import faulthandler
 import struct
@@ -138,15 +139,21 @@ class ANE:
 
   def debug(self, dat, mems=0):
     add = [0x30, 0x1d4, 0x220, 0x29c, 0x2f0, 0x30c, 0x32c]
+    lens = [244, 60, 108, 68, 12, 16, 24]
     ptr = 0x2b
     ddat = dat[0:0x28]
-    for a in add:
-      pm = dat[ptr]
+    for a, pm in zip(add, lens):
+      #assert pm == dat[ptr]
       ddat += b"\x00" * (a-len(ddat))
       ddat += dat[ptr+1:ptr+1+pm+4]
       ptr += pm+8
     ddat += b"\x00" * 8
-    return libane.ANE_RegDebug(0, create_string_buffer(ddat), mems).decode('utf-8')
+    ret = collections.OrderedDict()
+    for ln in libane.ANE_RegDebug(0, create_string_buffer(ddat), mems).decode('utf-8').strip().split("\n"):
+      lnn = ln.split(" = ")
+      if len(lnn) == 2:
+        ret[lnn[0]] = int(lnn[1])
+    return ret
 
   def filln(self, dat, nvdict, base=0x4000):
     for n,v in nvdict.items():