#! /usr/bin/env python3

# See LICENSE.SiFive for license details.
# See LICENSE.Berkeley for license details.

import sys
import math

use_latches = 0

class VerilogModuleGenerator(object):
  def __init__(self, name):
    self.name = name
    self.port_spec = []
    self.decl = []
    self.combinational = []
    self.sequential = []

  def __format_width(self, width):
    return "[{}:0] ".format(width-1) if width > 1 else ""

  def __format_depth(self, depth):
    return " [{}:0]".format(depth-1) if depth > 1 else ""

  def add_io(self, io_type, width, name):
    width_str = self.__format_width(width)
    # print(io_type, width_str, name)
    self.port_spec.append(f'{io_type} {width_str}{name}')

  def add_input(self, width, name):
    self.add_io("input", width, name)

  def add_output(self, width, name):
    self.add_io("output", width, name)

  def add_decl(self, decl_type, width, name, depth=1):
    width_str = self.__format_width(width)
    depth_str = self.__format_depth(depth)
    self.decl.append(f"{decl_type} {width_str}{name}{depth_str};")

  def add_decl_reg(self, width, name, depth=1):
    self.add_decl("reg", width, name, depth)

  def add_decl_ram(self, width, name, depth=1):
    width_str = self.__format_width(width)
    depth_str = " [{}:0]".format(depth-1)
    self.decl.append(f"reg {width_str}{name}{depth_str};")

  def add_decl_wire(self, width, name, depth=1):
    self.add_decl("wire", width, name, depth)

  def add_decl_line(self, line):
    self.decl.append(line)

  def add_sequential(self, line):
    self.sequential.append(line)

  def add_combinational(self, line):
    self.combinational.append(line)

  def generate(self, blackbox):
    body = "\
  %s\n\
  %s\n\
  %s\n" % ('\n  '.join(self.decl), '\n  '.join(self.sequential), '\n  '.join(self.combinational))

    s = "\nmodule %s(\n\
  %s\n\
);\n\
\n\
%s\
\n\
endmodule" % (self.name, ',\n  '.join(self.port_spec), body if not blackbox else blackbox)
    return s


class Reshaper(object):
  def __init__(self, before, after):
    # print(before, after)
    self.conf = before
    self.new_conf = after
    assert(self.conf[-1] == ['write', 'read'])
    assert(self.new_conf[-1] == ['mwrite', 'read'])

  def generate(self, mem):
    (name, width, depth, mask_gran, mask_seg, _) = self.conf
    (new_name, new_width, new_depth, new_mask_gran, new_mask_seg, _) = self.new_conf
    addr_bits = math.log2(depth)
    ways = new_width // width
    ways_bits = int(math.log2(ways))
    mem.add_decl_wire(new_width, "data_read")
    mem.add_decl_wire(new_width, "data_write")
    mem.add_combinational(f"assign data_write = ")
    sels = [f"{f'(write_way_index == {w}) ?' if w != ways-1 else ''} ({{{new_width-width}'h0, W0_data}} << {width*w})" for w in range(ways)]
    mem.add_combinational(":\n    ".join(sels) + ";")
    mem.add_decl_wire(ways_bits, "read_way_index")
    mem.add_combinational(f"assign read_way_index = R0_addr[{ways_bits-1}:0];")
    mem.add_decl_wire(ways_bits, "write_way_index")
    mem.add_combinational(f"assign write_way_index = W0_addr[{ways_bits-1}:0];")
    mem.add_combinational(f"{new_name} array (")
    mem.add_combinational(f"  .W0_clk(W0_clk),")
    mem.add_combinational(f"  .W0_addr(W0_addr[{new_width-1}:{ways_bits}]),")
    mem.add_combinational(f"  .W0_en(W0_en),")
    mem.add_combinational(f"  .W0_data(data_write),")
    mem.add_combinational(f"  .W0_mask({ways}'h1 << write_way_index),")
    mem.add_combinational(f"  .R0_clk(R0_clk),")
    mem.add_combinational(f"  .R0_addr(R0_addr[{new_width-1}:{ways_bits}]),")
    mem.add_combinational(f"  .R0_en(R0_en),")
    mem.add_combinational(f"  .R0_data(data_read)")
    mem.add_combinational(f");")
    mem.add_combinational(f"assign R0_data = ")
    sels = [f"{f'(read_way_index == {w}) ?' if w != ways-1 else ''} data_read[{width*(w+1)-1}:{width*w}]" for w in range(ways)]
    mem.add_combinational(":\n    ".join(sels) + ";")


class Spliter(object):
  def __init__(self, before, after):
    # print(before, after)
    self.conf = before
    self.new_conf = after
    assert(self.conf[-1] == ['mrw'])
    assert(self.new_conf[-1] == ['rw'])

  def generate(self, mem):
    (name, width, depth, mask_gran, mask_seg, _) = self.conf
    (new_name, new_width, new_depth, new_mask_gran, new_mask_seg, _) = self.new_conf
    assert(depth == new_depth)
    ways = width // new_width
    for i in range(ways):
      data_slice = f"[{new_width*(i+1)-1}:{new_width*i}]"
      mem.add_combinational(f"{new_name} array_{i} (")
      mem.add_combinational(f"  .RW0_clk(RW0_clk),")
      mem.add_combinational(f"  .RW0_addr(RW0_addr),")
      mem.add_combinational(f"  .RW0_en(RW0_en),")
      mem.add_combinational(f"  .RW0_wmode(RW0_wmode && RW0_wmask[{i}]),")
      mem.add_combinational(f"  .RW0_wdata(RW0_wdata{data_slice}),")
      mem.add_combinational(f"  .RW0_rdata(RW0_rdata{data_slice})")
      mem.add_combinational(f");")

class SRAM(object):
  def __init__(self, line):
    self.parse_line(line)
    self.prepare_module()

  def parse_line(self, line):
    name = ''
    width = 0
    depth = 0
    ports = ''
    mask_gran = 0
    tokens = line.split()
    i = 0
    for i in range(0, len(tokens), 2):
      s = tokens[i]
      if s == 'name':
        name = tokens[i+1]
      elif s == 'width':
        width = int(tokens[i+1])
        mask_gran = width # default setting
      elif s == 'depth':
        depth = int(tokens[i+1])
      elif s == 'ports':
        ports = tokens[i+1].split(',')
      elif s == 'mask_gran':
        mask_gran = int(tokens[i+1])
      else:
        sys.exit('%s: unknown argument %s' % (sys.argv[0], i))
    self.conf = (name, width, depth, mask_gran, width//mask_gran, ports)
  # return (name, width, depth, mask_gran, width//mask_gran, ports)

  def prepare_module(self):
    (name, width, depth, mask_gran, mask_seg, ports) = self.conf
    addr_width = max(math.ceil(math.log(depth)/math.log(2)),1)

    mem = VerilogModuleGenerator(name)
    readports = []
    writeports = []
    latchports = []
    rwports = []
    maskedports = {}

    for pid, ptype in enumerate(ports):
      if ptype[0:1] == 'm':
        ptype = ptype[1:]
        maskedports[pid] = pid

      if ptype == 'read':
        prefix = 'R%d_' % len(readports)
        mem.add_input(1, prefix + "clk")
        mem.add_input(addr_width, prefix + "addr")
        mem.add_input(1, prefix + "en")
        mem.add_output(width, prefix + "data")
        readports.append(pid)
      elif ptype == 'write':
        prefix = 'W%d_' % len(writeports)
        mem.add_input(1, prefix + "clk")
        mem.add_input(addr_width, prefix + "addr")
        mem.add_input(1, prefix + "en")
        mem.add_input(width, prefix + "data")
        if pid in maskedports:
          mem.add_input(mask_seg, prefix + "mask")
        if not use_latches or pid in maskedports:
          writeports.append(pid)
        else:
          latchports.append(pid)
      elif ptype == 'rw':
        prefix = 'RW%d_' % len(rwports)
        mem.add_input(1, prefix + "clk")
        mem.add_input(addr_width, prefix + "addr")
        mem.add_input(1, prefix + "en")
        mem.add_input(1, prefix + "wmode")
        if pid in maskedports:
          mem.add_input(mask_seg, prefix + "wmask")
        mem.add_input(width, prefix + "wdata")
        mem.add_output(width, prefix + "rdata")
        rwports.append(pid)
      else:
        sys.exit('%s: unknown port type %s' % (sys.argv[0], ptype))
    self.mem = mem
    self.ports_conf = (readports, writeports, latchports, rwports, maskedports)

  def generate(self, blackbox):
    (name, width, depth, mask_gran, mask_seg, ports) = self.conf
    addr_width = max(math.ceil(math.log(depth)/math.log(2)),1)
    mem, (readports, writeports, latchports, rwports, maskedports) = self.mem, self.ports_conf

    nr = len(readports)
    nw = len(writeports)
    nrw = len(rwports)

    def emit_read(idx, rw):
      prefix = ('RW%d_' if rw else 'R%d_') % idx
      data = ('%srdata' if rw else '%sdata') % prefix
      en = ('%sen && !%swmode' % (prefix, prefix)) if rw else ('%sen' % prefix)
      mem.add_decl_reg(1, f"reg_{prefix}ren")
      mem.add_decl_reg(addr_width, f"reg_{prefix}addr")
      mem.add_sequential(f"always @(posedge {prefix}clk)")
      mem.add_sequential(f"  reg_{prefix}ren <= {en};")
      mem.add_sequential(f"always @(posedge {prefix}clk)")
      mem.add_sequential(f"  if ({en}) reg_{prefix}addr <= {prefix}addr;")
      mem.add_combinational("`ifdef RANDOMIZE_GARBAGE_ASSIGN")
      mem.add_combinational(f"reg [{((width-1)//32+1)*32-1}:0] {prefix}random;")
      mem.add_combinational(f"`ifdef RANDOMIZE_MEM_INIT")
      mem.add_combinational(f"  initial begin")
      mem.add_combinational(f"    #`RANDOMIZE_DELAY begin end")
      mem.add_combinational('    %srandom = {%s};' % (prefix, ', '.join(['$random'] * ((width-1)//32+1))))
      mem.add_combinational('    reg_%sren = %srandom[0];' % (prefix, prefix))
      mem.add_combinational('  end')
      mem.add_combinational('`endif')
      mem.add_combinational('always @(posedge %sclk) %srandom <= {%s};' % (prefix, prefix, ', '.join(['$random'] * ((width-1)//32+1))))
      mem.add_combinational('assign %s = reg_%sren ? ram[reg_%saddr] : %srandom[%d:0];' % (data, prefix, prefix, prefix, width-1))
      mem.add_combinational('`else')
      mem.add_combinational('assign %s = ram[reg_%saddr];' % (data, prefix))
      mem.add_combinational('`endif')

    for idx in range(nr):
      emit_read(idx, False)

    for idx in range(nrw):
      emit_read(idx, True)

    for idx in range(len(latchports)):
      prefix = 'W%d_' % idx
      mem.add_decl_reg(addr_width, f"latch_{prefix}addr")
      mem.add_decl_reg(width, f"latch_{prefix}data")
      mem.add_decl_reg(1, f"latch_{prefix}en")
      mem.add_combinational('always @(*) begin')
      mem.add_combinational('  if (!%sclk && %sen) latch_%saddr <= %saddr;' % (prefix, prefix, prefix, prefix))
      mem.add_combinational('  if (!%sclk && %sen) latch_%sdata <= %sdata;' % (prefix, prefix, prefix, prefix))
      mem.add_combinational('  if (!%sclk) latch_%sen <= %sen;' % (prefix, prefix, prefix))
      mem.add_combinational('end')
      mem.add_combinational('always @(*)')
      mem.add_combinational('  if (%sclk && latch_%sen)' % (prefix, prefix))
      mem.add_combinational('    ram[latch_%saddr] <= latch_%sdata;' % (prefix, prefix))

    mem.add_decl_ram(width, "ram", depth)
    mem.add_decl_line('`ifdef RANDOMIZE_MEM_INIT')
    mem.add_decl_line('  integer initvar;')
    mem.add_decl_line('  initial begin')
    mem.add_decl_line('    #`RANDOMIZE_DELAY begin end')
    mem.add_decl_line('    for (initvar = 0; initvar < %d; initvar = initvar+1)' % depth)
    mem.add_decl_line('      ram[initvar] = {%d {$random}};' % ((width-1)//32+1))
    for idx in range(nr):
      prefix = 'R%d_' % idx
      mem.add_decl_line('    reg_%saddr = {%d {$random}};' % (prefix, ((addr_width-1)//32+1)))
    for idx in range(nrw):
      prefix = 'RW%d_' % idx
      mem.add_decl_line('    reg_%saddr = {%d {$random}};' % (prefix, ((addr_width-1)//32+1)))
    mem.add_decl_line('  end')
    mem.add_decl_line('`endif')

    mem.add_decl_line("integer i;")
    for idx in range(nw):
      prefix = 'W%d_' % idx
      pid = writeports[idx]
      mem.add_sequential('always @(posedge %sclk)' % prefix)
      mem.add_sequential("  if (%sen) begin" % prefix)
      for i in range(mask_seg):
        mask = ('if (%smask[%d]) ' % (prefix, i)) if pid in maskedports else ''
        ram_range = '%d:%d' % ((i+1)*mask_gran-1, i*mask_gran)
        mem.add_sequential("    %sram[%saddr][%s] <= %sdata[%s];" % (mask, prefix, ram_range, prefix, ram_range))
      mem.add_sequential("  end")
    for idx in range(nrw):
      pid = rwports[idx]
      prefix = 'RW%d_' % idx
      mem.add_sequential('always @(posedge %sclk)' % prefix)
      mem.add_sequential("  if (%sen && %swmode) begin" % (prefix, prefix))
      if mask_seg > 0:
        if mask_gran == 1: # If 1 bit mask, use & instead
          if pid in maskedports:
            mem.add_sequential("      ram[%saddr] <= (%swmask & %swdata) | (~%swmask & ram[%saddr]);" %(prefix, prefix, prefix, prefix, prefix))
          else:
            mem.add_sequential("      ram[%saddr] <= %swdata;" %(prefix, prefix))
        else:
          mem.add_sequential("    for (i=0;i<%d;i=i+1) begin" % mask_seg)
          if pid in maskedports:
            mem.add_sequential("      if (%swmask[i]) begin" % prefix)
            mem.add_sequential("        ram[%saddr][i*%d +: %d] <= %swdata[i*%d +: %d];" %(prefix, mask_gran, mask_gran, prefix, mask_gran, mask_gran))
            mem.add_sequential("      end")
          else:
            mem.add_sequential("      ram[%saddr][i*%d +: %d] <= %swdata[i*%d +: %d];" %(prefix, mask_gran, mask_gran, prefix, mask_gran, mask_gran))
          mem.add_sequential("    end")
      mem.add_sequential("  end")
    return mem.generate(blackbox)


class SRAM_TSMC28(SRAM):
  def __init__(self, line):
    super().__init__(line)
    self.sub_srams = []
    if self.__check_subsrams():
      print(line.strip())

  def __check_subsrams(self):
    need_split = self.__split()
    need_reshape  = self.__reshape()
    assert(not (need_split and need_reshape))
    return not need_split and not need_reshape

  def __split(self):
    (name, width, depth, mask_gran, mask_seg, ports) = self.conf
    '''if ports == ["mrw"] and mask_gran >= 32:
      new_conf = (name + "_sub", str(depth), str(mask_gran), "rw")
      line_field = ("name", "depth", "width", "ports")
      new_line = " ".join(map(lambda x: " ".join(x), zip(line_field, new_conf)))
      new_sram = SRAM_TSMC28(new_line)
      self.sub_srams.append(new_sram)
      reshaper = Spliter(self.conf, new_sram.conf)
      reshaper.generate(self.mem)
      return True'''
    return False

  def __reshape(self):
    (name, width, depth, mask_gran, mask_seg, ports) = self.conf
    if width == 2 and depth == 256:
      new_conf = (name + "_sub", "64", "8", "mwrite,read", "2")
      line_field = ("name", "depth", "width", "ports", "mask_gran")
      new_line = " ".join(map(lambda x: " ".join(x), zip(line_field, new_conf)))
      new_sram = SRAM_TSMC28(new_line)
      self.sub_srams.append(new_sram)
      reshaper = Reshaper(self.conf, new_sram.conf)
      reshaper.generate(self.mem)
      return True
    return False

  def __get_tsmc_lib(self):
    mem, (readports, writeports, latchports, rwports, maskedports) = self.mem, self.ports_conf
    blackbox = "// tsmc lib here\n"
    (name, width, depth, mask_gran, mask_seg, _) = self.conf
    nports = (len(readports), len(writeports), len(rwports))
    addr_width = max(math.ceil(math.log(depth)/math.log(2)),1)
    masked = len(maskedports) > 0
    # from tsmc28_sram import gen_tsmc_ram_1pw, gen_tsmc_ram_1pnw, gen_tsmc_ram_2pw, gen_tsmc_ram_2pnw
    # if nports == (1, 1, 0):
    #   if masked:
    #     blackbox = gen_tsmc_ram_2pw("TS6N28HPCPLVTA64X8M2F", width, mask_gran)
    #   else:
    #     blackbox = gen_tsmc_ram_2pnw("TS6N28HPCPLVTA64X14M2F")
    # elif nports == (0, 0, 1):
    #   if masked:
    #     blackbox = gen_tsmc_ram_1pw('TS1N28HPCPLVTB8192X64M8SW', width, mask_gran, addr_width)
    #   else:
    #     blackbox = gen_tsmc_ram_1pnw('TS5N28HPCPLVTA64X144M2F', width, addr_width)
    # else:
    #   blackbox = "// unknown tsmc lib type\n"
    return mem.generate(blackbox)

  def generate(self, blackbox, itself_only=False):
    if itself_only:
      # generate splits or reshapes
      if self.sub_srams:
        return self.mem.generate("")
      # use empty blackbox
      elif blackbox:
        return super().generate(" ")
      # insert tsmc libs
      else:
        return self.__get_tsmc_lib()
    else:
      s = self.generate(blackbox, True)
      for sram in self.sub_srams:
        s += sram.generate(blackbox)
      return s


def main(args):
  f = open(args.output_file, "w") if (args.output_file) else None
  conf_file = args.conf
  for line in open(conf_file):
    sram = SRAM(line)
    if args.tsmc28:
      sram = SRAM_TSMC28(line)
    else:
      sram = SRAM(line)
    if f is not None:
        f.write(sram.generate(args.blackbox))
    else:
        print(sram.generate(args.blackbox))


if __name__ == '__main__':
  import argparse
  parser = argparse.ArgumentParser(description='Memory generator for Rocket Chip')
  parser.add_argument('conf', metavar='.conf file')
  parser.add_argument('--tsmc28', action='store_true', help='use tsmc28 sram to generate module body')
  parser.add_argument('--blackbox', '-b', action='store_true', help='set to disable output of module body')
  #parser.add_argument('--use_latches', '-l', action='store_true', help='set to enable use of latches')
  parser.add_argument('--output_file', '-o', help='name of output file, default is stdout')
  args = parser.parse_args()
  #use_latches = args.use_latches
  main(args)