#!/usr/bin/env pypy3

import os
import sys
import string
import multiprocessing

def processwords(W):
  result = []
  prevword = ''
  parenstack = []
  commas = []
  for w in W:
    if w[0] in ('(','['):
      parenstack.append(prevword)
      commas.append(0)
    elif w[0] == ',':
      if len(commas) > 0:
        commas[-1] += 1
    elif w[0] in (')',']'):
      if len(parenstack) > 0:
        thisword = parenstack.pop()
        thiscommas = commas.pop()
        if thiscommas == 1:
          if thisword in ('Tyapp','Const','Var','Comb','Abs'):
            result.append([',',''])
            result.append(['_',''])
    result.append(w)
    prevword = w[0]
  return result

def process(x):
  W = [['','']]
  for c in x:
    if c.isalpha():
      if len(W[-1][1]) > 0 or (len(W[-1][0]) > 0 and not W[-1][0].isalpha()):
        W += [['','']]
      W[-1][0] += c
    elif c.isspace():
      W[-1][1] += c
    else:
      if W[-1] != ['','']:
        W += [['','']]
      W[-1][0] += c
  W = processwords(W)
  return ''.join(''.join(w) for w in W)

def process_fusion(x):
  x = process(x)
  result = ''
  for line in x.splitlines():
    line = line.replace('Tyapp of string *  hol_type list','Tyapp of string *  hol_type list * int')
    line = line.replace('Var of string * hol_type','Var of string * hol_type * int')
    line = line.replace('Const of string * hol_type','Const of string * hol_type * int')
    line = line.replace('Comb of term * term','Comb of term * term * int')
    line = line.replace('Abs of term * term','Abs of term * term * int')

    line = line.replace('Tyapp(tyop,args,_)','unchecked_mk_type(tyop,args)')
    line = line.replace("Tyapp(tycon,args',_)","unchecked_mk_type(tycon,args')")
    line = line.replace('Tyapp("bool",[],_)','unchecked_mk_type("bool",[])')
    line = line.replace('["=",Tyapp("fun",[aty;Tyapp("fun",[aty;bool_ty],_)],_)]','["=",unchecked_mk_type("fun",[aty;unchecked_mk_type("fun",[aty;bool_ty])])]')
    line = line.replace('Tyapp("fun",[ty;type_of t],_)','unchecked_mk_type("fun",[ty;type_of t])')
    line = line.replace('Var(v,ty,_)','let h = Hashtbl.hash(3,v,hash_type ty) in Var(v,ty,h)')
    line = line.replace('Const(name,type_subst theta uty,_)','unchecked_mk_const(name,type_subst theta uty)')
    line = line.replace('Abs(bvar,bod,_)','unchecked_mk_abs(bvar,bod)')
    line = line.replace('Comb(f,a,_)','unchecked_mk_comb(f,a)')
    line = line.replace('Var(s^"\'",ty,_)','mk_var(s^"\'",ty)')
    line = line.replace("Comb(s',t',_)","unchecked_mk_comb(s',t')")
    line = line.replace("Abs(v',vsubst ((v',v)::ilist') s,_)","unchecked_mk_abs(v',vsubst ((v',v)::ilist') s)")
    line = line.replace("Abs(v,s',_)","unchecked_mk_abs(v,s')")
    line = line.replace("Var(n,ty',_)","mk_var(n,ty')")
    line = line.replace("Const(c,ty',_)","unchecked_mk_const(c,ty')")
    line = line.replace("Comb(f',x',_)","unchecked_mk_comb(f',x')")
    line = line.replace("Abs(y',t',_)","unchecked_mk_abs(y',t')")
    line = line.replace("Var(fst(dest_var y''),snd(dest_var y),_)","mk_var(fst(dest_var y''),snd(dest_var y))")
    line = line.replace("(Abs(z,vsubst[z,y] t,_))","(unchecked_mk_abs(z,vsubst[z,y] t))")
    line = line.replace('Comb(Comb(Const("=",Tyapp("fun",[ty;Tyapp("fun",[ty;bool_ty],_)],_),_),l,_),r,_)','unchecked_mk_comb(unchecked_mk_comb(unchecked_mk_const("=",unchecked_mk_type("fun",[ty;unchecked_mk_type("fun",[ty;bool_ty])])),l),r)')
    line = line.replace('Comb(eql,r,_)','unchecked_mk_comb(eql,r)')
    line = line.replace('(Comb(l1,l2,_))','(unchecked_mk_comb(l1,l2))')
    line = line.replace('(Comb(r1,r2,_))','(unchecked_mk_comb(r1,r2))')
    line = line.replace('(Abs(v,l,_))','(unchecked_mk_abs(v,l))')
    line = line.replace('(Abs(v,r,_))','(unchecked_mk_abs(v,r))')
    line = line.replace('; Const(cname,ty,_)','; unchecked_mk_const(cname,ty)')
    line = line.replace('= Tyapp(tyname,tyvars,_)','= unchecked_mk_type(tyname,tyvars)')
    line = line.replace('absty = Tyapp("fun",[rty;aty],_) and repty = Tyapp("fun",[aty;rty],_)','absty = unchecked_mk_type("fun",[rty;aty]) and repty = unchecked_mk_type("fun",[aty;rty])')
    line = line.replace('; Const(absname,absty,_)','; unchecked_mk_const(absname,absty)')
    line = line.replace('; Const(repname,repty,_)','; unchecked_mk_const(repname,repty)')
    line = line.replace('a = Var("a",aty,_) and r = Var("r",rty,_)','a = mk_var("a",aty) and r = mk_var("r",rty)')
    line = line.replace('(Comb(abs,mk_comb(rep,a),_))','(unchecked_mk_comb(abs,mk_comb(rep,a)))')
    line = line.replace('(Comb(P,r,_))','(unchecked_mk_comb(P,r))')

    if line.strip() == 'let mk_type(tyop,args) =':
      result += r'''
  let unchecked_mk_type(tyop,args) =
    let h = Hashtbl.hash(1,tyop,map hash_type args) in
    Tyapp(tyop,args,h)
'''
    if line.strip() == 'let mk_const(name,theta) =':
      result += r'''
  let unchecked_mk_const(name,ty) =
    let h = Hashtbl.hash(4,name,hash_type ty) in
    Const(name,ty,h)
'''
    if line.strip() == 'let mk_abs(bvar,bod) =':
      result += r'''
  let unchecked_mk_abs(bvar,bod) =
    let h = Hashtbl.hash(5,hash_term bvar,hash_term bod) in
    Abs(bvar,bod,h)
'''
    if line.strip() == 'let mk_comb(f,a) =':
      result += r'''
  let unchecked_mk_comb(f,a) =
    let h = Hashtbl.hash(hash_term f,hash_term a) in
    Comb(f,a,h)
'''

    result += line+'\n'

    if line.strip() == 'type thm':
      result += r'''
      val hash_type: hol_type -> int
      val hash_term: term -> int
'''
    if line.strip() == 'type thm = Sequent of (term list * term)':
      result += r'''
  let hash_type =
    function
      Tyvar v -> Hashtbl.hash(0,v)
    | Tyapp(_,_,h) -> h

  let hash_term =
    function
      Var(_,_,h) -> h
    | Const(_,_,h) -> h
    | Comb(_,_,h) -> h
    | Abs(_,_,h) -> h
'''

  return result

def process_holtrace(x):
  x = process(x)

  result = ''
  for line in x.splitlines():

    if line.strip() == '(* ===== types *)':
      result += r'''
  module Hashtbl_type = Hashtbl.Make(struct
    type t = hol_type
    let equal ty0 ty1 = (hash_type ty0 = hash_type ty1) && (ty0 = ty1)
    let hash = hash_type
  end)
'''
    if line.strip() == '(* ===== terms *)':
      result += r'''
  module Hashtbl_term = Hashtbl.Make(struct
    type t = term
    let equal tm0 tm1 = (hash_term tm0 = hash_term tm1) && (tm0 = tm1)
    let hash = hash_term
  end)
'''
    if line.strip() == '(* ===== theorems *)':
      result += r'''
  let hash_hc(h,c) = Hashtbl.hash(map hash_term (c::h))

  module Hashtbl_thm = Hashtbl.Make(struct
    type t = term list * term
    let equal hc0 hc1 = (hash_hc hc0 = hash_hc hc1) && (hc0 = hc1)
    let hash = hash_hc
  end)
'''

    if 'Hashtbl.' in line:
      if 'type2id' in line:
        line = line.replace('Hashtbl.','Hashtbl_type.')
      elif 'term2id' in line:
        line = line.replace('Hashtbl.','Hashtbl_term.')
      elif 'thmhc2id' in line:
        line = line.replace('Hashtbl.','Hashtbl_thm.')

    result += line+'\n'

  return result

# ===== main

if len(sys.argv) <= 1:
  raise Exception('please specify target directory')
dir = sys.argv[1]
if not os.path.exists(f'{dir}/holtrace.ml'):
  raise Exception('missing holtrace.ml in target directory')
if os.path.exists(f'{dir}/PATCHED-hashtt'):
  raise Exception('target directory patched already')

with open(f'{dir}/PATCHED-hashtt','w') as f:
  pass

def todo():
  for (subdir,dirnames,filenames) in os.walk(dir):
    for fn in filenames:
      if fn.endswith('.ml'):
        yield f'{subdir}/{fn}'

def doit(fn):
  try:
    with open(fn) as f: x = f.read()
  except:
    return
  if fn.endswith('/fusion.ml'):
    y = process_fusion(x)
  elif fn.endswith('/holtrace.ml'):
    y = process_holtrace(x)
  else:
    y = process(x)
  if y == x: return
  mode = os.stat(fn).st_mode
  os.rename(fn,fn+'.hashtt-bak')
  with open(fn,'w') as f: f.write(y)
  os.chmod(fn,mode)

with multiprocessing.Pool(16) as p:
  for result in p.imap_unordered(doit,todo(),1):
    pass