1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4#
5# Copyright 2017, Data61
6# Commonwealth Scientific and Industrial Research Organisation (CSIRO)
7# ABN 41 687 119 230.
8#
9# This software may be distributed and modified according to the terms of
10# the BSD 2-Clause license. Note that NO WARRANTY is provided.
11# See "LICENSE_BSD2.txt" for details.
12#
13# @TAG(DATA61_BSD)
14#
15
16'''
17An example of how to use the AST traversal functionality.
18'''
19
20from __future__ import absolute_import, division, print_function, \
21    unicode_literals
22
23import os, sys
24sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
25import camkes.ast as ast
26import camkes.parser as parser
27
28def basic_visit(parent, node, ignored):
29    sys.stdout.write(' %s\n' % type(node))
30    return ast.TRAVERSAL_RECURSE
31
32def code_gen_enter(parent, node, state):
33    if isinstance(node, ast.Method):
34        if node.return_type is not None:
35            sys.stdout.write(node.return_type)
36        else:
37            sys.stdout.write('void ')
38        sys.stdout.write('%s(' % node.name)
39        state['infunction'] = True
40        state['firstparameter'] = True
41        return ast.TRAVERSAL_RECURSE
42    elif isinstance(node, ast.Parameter) and state.get('infunction', False):
43        if not state.get('firstparameter', True):
44            sys.stdout.write(', ')
45        sys.stdout.write('%s %s' % (node.type, node.name))
46        state['firstparameter'] = False
47        return ast.TRAVERSAL_CONTINUE
48    return ast.TRAVERSAL_RECURSE
49
50def code_gen_exit(parent, node, ignored):
51    if isinstance(node, ast.Method):
52        sys.stdout.write(') {\n  /* hello world */\n}\n')
53    return ast.TRAVERSAL_RECURSE
54
55def code_constructor(parent, node, state):
56    if not isinstance(node, ast.Parameter) and not isinstance(node, ast.Type):
57        state['infunction'] = None
58
59    if isinstance(node, ast.Method):
60        if node.name not in state['functions']:
61            state['functions'][node.name] = []
62            state['infunction'] = node.name
63        return ast.TRAVERSAL_RECURSE
64    elif isinstance(node, ast.Parameter) or isinstance(node, ast.Type) and state['infunction'] is not None:
65        state['functions'][state['infunction']].append(node)
66        return ast.TRAVERSAL_CONTINUE
67    return ast.TRAVERSAL_RECURSE
68
69def main():
70    if len(sys.argv) != 2:
71        sys.stderr.write('Usage: %s inputfile\n' % sys.argv[0])
72        return -1
73
74    with open(sys.argv[1], 'rt') as f:
75        s = f.read()
76
77    a = parser.parse_to_ast(s)
78
79    sys.stdout.write('Traverse the AST and print the types of nodes:\n')
80    ast.traverse(a, basic_visit, None, None)
81
82    sys.stdout.write('\nNow let\'s try some basic online code generation:\n')
83    ast.traverse(a, code_gen_enter, code_gen_exit, {})
84
85    sys.stdout.write('\nHow about the same offline:\n')
86    state = {
87        'functions':{},
88        'infunction':None,
89    }
90    ast.traverse(a, code_constructor, None, state)
91    for k, v in state['functions'].items():
92        if v[0] is not None:
93            sys.stdout.write(v[0])
94        else:
95            sys.stdout.write('void ')
96        sys.stdout.write('%(name)s(%(params)s) {\n  /* hello world */\n}\n' % {
97            'name':k,
98            'params':', '.join(map(lambda x: '%s %s' % (x.type, x.name), v[1:])),
99        })
100
101    return 0
102
103if __name__ == '__main__':
104    sys.exit(main())
105