Prolog: use A,B,C,… for normalized variable names

This is somewhat more readable than A0,A1,A2,….

1 files changed, 12 insertions, 3 deletions

diff --git a/prolog/util.py b/prolog/util.py
index e682edb..83f891f 100644
--- a/prolog/util.py
+++ b/prolog/util.py
@@ -16,6 +16,8 @@
 
 from collections import namedtuple
 from collections.abc import Iterable
+import string
+
 from nltk import Tree
 
 # Stores a token's type and value, and optionally the position of the first
@@ -78,6 +80,13 @@ def stringify(obj):
             return ''.join([stringify(child) for child in obj]) + '\n'
         return ''.join([stringify(child) for child in obj])
 
+# Return a canonical name for the [n]th variable in scope.
+def canonical_varname(n):
+    names = string.ascii_uppercase
+    if n < len(names):
+        return names[n]
+    return 'X{}'.format(n)
+
 # Rename variables in [tokens] to A0, A1, A2,… in order of appearance.
 def rename_vars_list(tokens, names=None):
     if names is None:
@@ -90,7 +99,7 @@ def rename_vars_list(tokens, names=None):
         if t.type == 'VARIABLE' and t.val != '_':
             cur_name = t.val
             if cur_name not in names:
-                names[cur_name] = 'A{}'.format(next_id)
+                names[cur_name] = canonical_varname(next_id)
                 next_id += 1
             tokens[i] = t.clone(val=names[cur_name])
     return tokens
@@ -115,7 +124,7 @@ def rename_vars_ast(root, fixed_names=None):
             if node.type == 'VARIABLE':
                 token = node
                 if token.val.startswith('_'):
-                    new_node = token.clone(val='A{}'.format(next_id))
+                    new_node = token.clone(val=canonical_varname(next_id))
                     next_id += 1
                 else:
                     cur_name = token.val
@@ -123,7 +132,7 @@ def rename_vars_ast(root, fixed_names=None):
                         new_name = fixed_names[cur_name]
                     else:
                         if cur_name not in names:
-                            names[cur_name] = 'A{}'.format(next_id)
+                            names[cur_name] = canonical_varname(next_id)
                             next_id += 1
                         new_name = names[cur_name]
                     new_node = token.clone(val=new_name)