def test_SerializeInstructionsInProgramGraph_single_function():
proto = program_graph_pb2.ProgramGraph(
node=[
node_pb2.Node(type=node_pb2.Node.INSTRUCTION),
node_pb2.Node(type=node_pb2.Node.INSTRUCTION),
node_pb2.Node(type=node_pb2.Node.INSTRUCTION),
],
edge=[
edge_pb2.Edge(flow=edge_pb2.Edge.CALL, source=0, target=1),
edge_pb2.Edge(flow=edge_pb2.Edge.CONTROL, source=1, target=2),
],
)
n = graph_serializer.SerializeInstructionsInProgramGraph(proto,
max_nodes=1000)
assert n == [1, 2]
def test_GraphTuple_one_graph():
graph = program_graph_pb2.ProgramGraph(
node=[node_pb2.Node(),],
edge=[
edge_pb2.Edge(source=0, target=1,),
edge_pb2.Edge(source=0, target=2, position=1,),
edge_pb2.Edge(source=1, target=0, position=10, flow=edge_pb2.Edge.CALL,),
],
)
builder = GraphTupleBuilder()
builder.AddProgramGraph(graph)
gt = builder.Build()
assert np.array_equal(gt.adjacencies[edge_pb2.Edge.CONTROL], [(0, 1), (0, 2)])
assert np.array_equal(
gt.adjacencies[edge_pb2.Edge.DATA], np.zeros((0, 2), dtype=np.int32)
)
assert np.array_equal(gt.adjacencies[edge_pb2.Edge.CALL], [(1, 0)])
assert np.array_equal(gt.edge_positions[edge_pb2.Edge.CONTROL], [0, 1])
assert np.array_equal(gt.edge_positions[edge_pb2.Edge.DATA], [])
assert np.array_equal(gt.edge_positions[edge_pb2.Edge.CALL], [10])