from typing import Union, List, NamedTuple

# ASDL's seven builtin types are:
# - identifier, int, string, bytes, object, singleton, constant
# - singleton: None, True or False
# - constant can be None, whereas None means "no value" for object.
#
# Resources for ASDL and ASTs:
# cf. https://www.usenix.org/legacy/publications/library/proceedings/dsl97/full_papers/wang/wang.pdf
# cf. https://docs.python.org/3/library/ast.html#module-ast
# cf. https://eli.thegreenplace.net/2014/06/04/using-asdl-to-describe-asts-in-compilers
# cf. https://www.oilshell.org/blog/2016/12/11.html

asdl_spec = '''\
start = Program(stmt* procs, expr* calls)

stmt = Procedure(identifier name, identifier* params, bool is_test, stmt body)
     | Block(int blocknum, stmt* stmts)
     | Assign(lvalue target, expr value)
     | If(expr cond, stmt action)
     | Loop(int times, bool fixed, body stmt)
     | AbortLoop(int blocknum)
     | QuitBlock(int blocknum)

expr = BinOp(expr value1, str op, expr value2)
     | Number(int x)
     | Bool(bool x)
     | Id(identifier name)
     | Call(identifier name, expr* args)
     | lvalue

lvalue = Output(bool is_bool)
       | Cell(int i)
'''

identifier = str
stmt = Union['Procedure', 'Block','Assign', 'If', 'Loop', 'AbortLoop', 'QuitBlock']
expr = Union['BinOp', 'Number', 'Bool', 'Id', 'Call', 'Output', 'Cell']

class Program(NamedTuple):
    procs: List[stmt]
    calls: List[expr] = None

class Procedure(NamedTuple):
    name: identifier
    params: List[identifier]
    is_test: bool
    body: stmt

class Block(NamedTuple):
    blocknum: int
    stmts: List[stmt]

class Assign(NamedTuple):
    name: 'Union[Cell, Output]'
    value: expr

class If(NamedTuple):
    cond: expr
    action: stmt

class Loop(NamedTuple):
    times: int
    fixed: bool
    body: stmt

class AbortLoop(NamedTuple):
    blocknum: int

class QuitBlock(NamedTuple):
    blocknum: int

class BinOp(NamedTuple):
    value1: expr
    op: str
    value2: expr

class Number(NamedTuple):
    x: int

class Bool(NamedTuple):
    x: bool

class Id(NamedTuple):
    name: identifier

class Call(NamedTuple):
    name: identifier
    args: List[expr]

class Output(NamedTuple):
    is_bool: bool

class Cell(NamedTuple):
    i: int

#######################################################################

atomic = {identifier, int, str, bytes, type(None), bool}

constructors = {Program, Procedure, Block, Assign, If, Loop, AbortLoop,
                QuitBlock, BinOp, Number, Bool, Id, Call, Output, Cell}

def children(node):
    'All direct child nodes that are themselves constructors'
    for subnode in node:
        # Only one level flattening of lists in needed in ASDL
        for ssn in (subnode if type(subnode) is list else [subnode]):
            if type(ssn) in constructors:
                yield ssn

indent = 3

def display(node, pad=''):
    morepad = pad + ' ' * indent
    typ = type(node)
    if typ in atomic:
        return f'{node!r}'
    elif typ in constructors:
        if len(node) == 1:
            return f'{node!r}'
        result = [f'{type(node).__name__}(']
        rows = [f'{morepad}{field_name} = {display(value, morepad)}'
            for field_name, value in zip(node._fields, node)]
        result += [',\n'.join(rows)]
        result += [f'{pad})']
        return '\n'.join(result)
    elif typ is list:
        if not node:
            return '[]'
        result = ['[']
        rows = [f'{morepad}{display(value, morepad)}'
                for value in node]
        result += [',\n'.join(rows)]
        result += [f'{pad}]']
        return '\n'.join(result)
    else:
        raise TypeError('Unknown node type', node)