diff --git a/cozy/main.py b/cozy/main.py index e4d16e5d..8983d81c 100755 --- a/cozy/main.py +++ b/cozy/main.py @@ -101,7 +101,12 @@ def run(): start = datetime.datetime.now() - if not args.simple: + if args.simple: + if args.save: + with open(args.save, "wb") as f: + pickle.dump(ast, f) + print("Saved implementation to file {}".format(args.save)) + else: callback = None server = None @@ -146,16 +151,12 @@ def callback(impl): ast, timeout = datetime.timedelta(seconds=args.timeout), progress_callback = callback, - improve_count=improve_count) + improve_count=improve_count, + dump_synthesized_in_file=args.save) if server is not None: server.join() - if args.save: - with open(args.save, "wb") as f: - pickle.dump(ast, f) - print("Saved implementation to file {}".format(args.save)) - print("Generating IR...") code = ast.code diff --git a/cozy/synthesis/high_level_interface.py b/cozy/synthesis/high_level_interface.py index 1edb6a5b..bc28d1c2 100644 --- a/cozy/synthesis/high_level_interface.py +++ b/cozy/synthesis/high_level_interface.py @@ -10,6 +10,7 @@ from typing import Callable, Any import sys import os +import pickle from queue import Empty from multiprocessing import Value @@ -99,7 +100,8 @@ def improve_implementation( impl : Implementation, timeout : datetime.timedelta = datetime.timedelta(seconds=60), progress_callback : Callable[[Implementation], Any] = None, - improve_count : Value = None) -> Implementation: + improve_count : Value = None, + dump_synthesized_in_file: str = None) -> Implementation: """Improve an implementation. This function tries to synthesize a better version of the given @@ -108,6 +110,10 @@ def improve_implementation( If provided, progress_callback will be called whenever a better implementation is found. It will be given the better implementation, which it should not modify or cache. + + If provided, the synthesized implementation will be dumped to dump_synthesized_in_file + before cleaning up the running threads when the loop terminates (because of time-outs etc.). + This is useful when thread that invokes Z3 is not responsive to cleanup. """ start_time = datetime.datetime.now() @@ -241,6 +247,11 @@ def index_of(l, p): else: print(" (skipped; {} was aleady cleaned up)".format(q.name)) + if dump_synthesized_in_file is not None: + with open(dump_synthesized_in_file, "wb") as f: + pickle.dump(impl, f) + print("Dumped implementation to file {}".format(dump_synthesized_in_file)) + # stop jobs print("Stopping jobs") stop_jobs(list(improvement_jobs))