<gh_stars>100-1000
#!/usr/bin/env python3
# Copyright 2020-2021 Efabless Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Original code by <NAME>, Translated to Python by Donn

import io
import re
import math
import click
import textwrap

@click.command()
@click.option("--fanout", required=True)
@click.option("--clk-net", required=True)
@click.option("--root-clkbuf", required=True, help="Clock tree root buffer type")
@click.option("--clkbuf", required=True, help="Clock tree branching buffer type")
@click.option("--clkbuf-input-pin", required=True, help="Name of input pin in clock buffers")
@click.option("--clkbuf-output-pin", required=True, help="Name of output pin in clock buffers")
@click.option("--clk-port", required=True, help="Name of clock pin in storage elements")
@click.option("-o", "--output", required=True, help="Name of output netlist")
@click.argument("input_netlist")
def cli(fanout, clk_net, clkbuf, root_clkbuf, clkbuf_input_pin, clkbuf_output_pin, clk_port, output, input_netlist):
    """
    Shalan's Simple Clock Tree Synthesizer

    Made to work on Yosys netlists.
    """
    fanout = int(fanout)

    netlist_str = open(input_netlist).read()
    clock_port_rx = re.compile(rf"{re.escape(clk_port)}\(\s*{re.escape(clk_net)}\s*\)")
    leaves = len(clock_port_rx.findall(netlist_str))
    levels = math.ceil(math.log(leaves, fanout))

    buffers = [0] * (levels + 1)
    buffers[0] = leaves

    verilog_wires = []
    verilog_cells = []

    cell_count = 0
    for level in range(1, levels):
        level_before = level - 1
        instance = f"_CTS_buf_{level}_"
        input_wire = f"clk_{level}_"
        output_wire = f"clk_{level_before}_"

        # obviously translated c-style for loop
        i = 0
        while i < leaves:
            next_level_fanout = (fanout ** (level + 1))
            ii = i//next_level_fanout * next_level_fanout
            cell_name = f"{instance}{i}"
            verilog_cells.append(textwrap.dedent(f"""
            {clkbuf} {cell_name} (
            \t.{clkbuf_input_pin}({input_wire}{ii}),
            \t.{clkbuf_output_pin}({output_wire}{i})
            );"""))
            verilog_wires.append(f"wire {output_wire}{i};")
            buffers[level] += 1
            cell_count += 1

            i += fanout ** level

    root_net = f"clk_{levels - 1}_0"
    verilog_cells.append(textwrap.dedent(f"""
    {root_clkbuf} _CTS_root(
    \t.{clkbuf_input_pin}({clk_net}),
    \t.{clkbuf_output_pin}({root_net})
    );"""))
    verilog_wires.append(f"wire clk_{levels - 1}_0;\n")

    with io.StringIO() as sio:
        state = 0
        ff_count = 0
        for line in netlist_str.split("\n"):
            if state == 0:
                if "wire" in line:
                    state = 1
                print(line, file=sio)
            elif state == 1:
                if not "wire" in line and not "input" in line and not "output" in line:
                    state = 2
                    print("\n// CTS added wires:", file=sio)
                    print("\n".join(verilog_wires), file=sio)
                    print("\n// CTS added buffers:", file=sio)
                    print("\n".join(verilog_cells), file=sio)
                    print(line, file=sio)
                else:
                    print(line, file=sio)
            elif state == 2:
                clk_port = clock_port_rx.search(line)
                if clk_port is not None:
                    clk_wire_name = f"clk_0_{ff_count // fanout * fanout}"
                    line_replaced = re.sub(re.escape(clk_net), clk_wire_name, line)
                    ff_count += 1
                    print(line_replaced, file=sio)
                else:
                    print(line, file=sio)
        with open(output, 'w') as f:
            f.write(sio.getvalue())

if __name__ == "__main__":
    cli()