Count the number of bytes a program uses

Author: dwayne

herbert/counter.py

@@ -0,0 +1,59 @@
+import itertools
+
+
+def count_bytes(parse_tree):
+    assert parse_tree.data == 'h'
+
+    *pdefs, main = parse_tree.children
+
+    return sum(itertools.chain(
+        map(_count_pdef, pdefs),
+        map(_count_seq, main.children)
+    ))
+
+
+def _count_pdef(pdef):
+    assert pdef.data == 'pdef'
+
+    try:
+        name, params, body = pdef.children
+        params = params.children
+    except:
+        name, body = pdef.children
+        params = []
+
+    return 1 + len(params) + sum(map(_count_seq, body.children))
+
+
+def _count_seq(x):
+    if hasattr(x, 'type'):
+        assert x.type == 'COMMAND' or x.type == 'PARAM'
+        return 1
+    else:
+        assert x.data == 'pcall'
+
+        try:
+            name, args = x.children
+            args = args.children
+        except ValueError:
+            name = x.children[0]
+            args = []
+
+        return 1 + sum(map(_count_arg, args))
+
+
+def _count_arg(arg):
+    if arg.data == 'var':
+        return 1
+
+    if arg.data == 'sexpr':
+        return sum(map(_count_seq, arg.children))
+
+    assert arg.data == 'expr'
+
+    l = len(arg.children)
+    if l % 2 == 1:
+        l += 1
+    l //= 2
+
+    return l

tests/test_counter.py

@@ -0,0 +1,35 @@
+import unittest
+
+from herbert.counter import count_bytes
+from herbert.parser import parse
+
+
+def count(program):
+    return count_bytes(parse(program))
+
+
+class CountBytesTestCase(unittest.TestCase):
+    def test_example1(self):
+        program = 'slr'
+
+        self.assertEqual(count(program), 3)
+
+    def test_example2(self):
+        program = 'a:sa\na'
+
+        self.assertEqual(count(program), 4)
+
+    def test_example3(self):
+        program = 'a(A):sa(A-1)\na(4)'
+
+        self.assertEqual(count(program), 8)
+
+    def test_example4(self):
+        program = 'a(A,B,C):f(B)Ca(A-1,B,C)\nf(A):sf(A-1)\na(4,5,rslsr)'
+
+        self.assertEqual(count(program), 26)
+
+    def test_example5(self):
+        program = 'f(A,B,C,D,E,F):sAf(AA,B,F-C+D+2,-D+1,E-1,sFl)\nf(sl,5,-1,100,10,r)'
+
+        self.assertEqual(count(program), 32)