From e294b241f8a9c35b99853b32d7c78e3cd8bc3a0d Mon Sep 17 00:00:00 2001
From: Mohamed Gaber <donn@efabless.com>
Date: Tue, 19 Nov 2024 22:25:19 +0200
Subject: [PATCH] add --initial-state-element-override/-e option

---
 openlane/__main__.py  | 52 +++++++++++++++++++++++++++++--------------
 openlane/flows/cli.py | 12 +++++++++-
 2 files changed, 46 insertions(+), 18 deletions(-)

diff --git a/openlane/__main__.py b/openlane/__main__.py
index 86780f0e5..a1dcef55b 100644
--- a/openlane/__main__.py
+++ b/openlane/__main__.py
@@ -23,12 +23,7 @@
 from functools import partial
 from typing import Any, Dict, Sequence, Tuple, Type, Optional, List, Union
 
-from click import (
-    Parameter,
-    Context,
-    Path,
-    pass_context,
-)
+import click
 from cloup import (
     option,
     option_group,
@@ -40,7 +35,7 @@
 
 
 from .__version__ import __version__
-from .state import State
+from .state import State, DesignFormat
 from .logging import (
     debug,
     err,
@@ -56,7 +51,7 @@
 
 
 def run(
-    ctx: Context,
+    ctx: click.Context,
     flow_name: Optional[str],
     pdk_root: Optional[str],
     pdk: str,
@@ -73,6 +68,7 @@ def run(
     config_override_strings: List[str],
     _force_run_dir: Optional[str],
     design_dir: Optional[str],
+    initial_state_element_override: Sequence[str],
     view_save_path: Optional[str] = None,
     ef_view_save_path: Optional[str] = None,
 ):
@@ -97,6 +93,27 @@ def run(
         if flow_description is None:
             flow_description = "Classic"
 
+        if len(initial_state_element_override):
+            if with_initial_state is None:
+                with_initial_state = State()
+            overrides = {}
+            for element in initial_state_element_override:
+                element_split = element.split("=", maxsplit=1)
+                if len(element_split) < 2:
+                    err(f"Invalid initial state element override: '{element}'.")
+                    ctx.exit(1)
+                df_id, path = element_split
+                design_format = DesignFormat.by_id(df_id)
+                if design_format is None:
+                    err(f"Invalid design format ID: '{df_id}'.")
+                    ctx.exit(1)
+                overrides[design_format] = common.Path(path)
+
+            with_initial_state = with_initial_state.__class__(
+                with_initial_state,
+                overrides=overrides,
+            )
+
         TargetFlow: Type[Flow]
 
         if isinstance(flow_description, str):
@@ -170,7 +187,7 @@ def run(
         flow._save_snapshot_ef(evsp)
 
 
-def print_version(ctx: Context, param: Parameter, value: bool):
+def print_version(ctx: click.Context, param: click.Parameter, value: bool):
     if not value:
         return
 
@@ -198,8 +215,8 @@ def print_version(ctx: Context, param: Parameter, value: bool):
 
 
 def print_bare_version(
-    ctx: Context,
-    param: Parameter,
+    ctx: click.Context,
+    param: click.Parameter,
     value: bool,
 ):
     if not value:
@@ -209,7 +226,7 @@ def print_bare_version(
 
 
 def run_included_example(
-    ctx: Context,
+    ctx: click.Context,
     smoke_test: bool,
     example: Optional[str],
     **kwargs,
@@ -285,8 +302,8 @@ def run_included_example(
 
 
 def cli_in_container(
-    ctx: Context,
-    param: Parameter,
+    ctx: click.Context,
+    param: click.Parameter,
     value: bool,
 ):
     if not value:
@@ -331,14 +348,14 @@ def cli_in_container(
     o(
         "--save-views-to",
         "view_save_path",
-        type=Path(file_okay=False, dir_okay=True),
+        type=click.Path(file_okay=False, dir_okay=True),
         default=None,
         help="A directory to copy the final views to, where each format is saved under a directory named after the corner ID (much like the 'final' directory after running a flow.)",
     ),
     o(
         "--ef-save-views-to",
         "ef_view_save_path",
-        type=Path(file_okay=False, dir_okay=True),
+        type=click.Path(file_okay=False, dir_okay=True),
         default=None,
         help="A directory to copy the final views to in the Efabless format, compatible with Caravel User Project.",
     ),
@@ -401,8 +418,9 @@ def cli_in_container(
     _enable_debug_flags=True,
     sequential_flow_reproducible=True,
     enable_overwrite_flag=True,
+    enable_initial_state_element=True,
 )
-@pass_context
+@click.pass_context
 def cli(ctx, /, **kwargs):
     """
     Runs an OpenLane flow via the commandline using a design configuration
diff --git a/openlane/flows/cli.py b/openlane/flows/cli.py
index 09e1b8d0d..6378c83d2 100644
--- a/openlane/flows/cli.py
+++ b/openlane/flows/cli.py
@@ -139,6 +139,7 @@ def cloup_flow_opts(
     volare_pdk_override: Optional[str] = None,
     _enable_debug_flags: bool = False,
     enable_overwrite_flag: bool = False,
+    enable_initial_state_element: bool = False,
 ) -> Decorator:
     """
     Creates a wrapper that appends a number of OpenLane flow-related flags to a
@@ -214,7 +215,7 @@ def decorate(f):
                 ),
                 default=None,
                 callback=initial_state_cb,
-                help="Use this JSON file as an initial state. If this is not specified, the latest `state_out.json` of the run directory will be used if available.",
+                help="Use this JSON file as an initial state. If this is not specified, the latest `state_out.json` of the run directory will be used. If none exist, an empty initial state is created.",
             )(f)
             f = o(
                 "--design-dir",
@@ -381,6 +382,15 @@ def decorate(f):
                 callback=set_worker_count_cb,
                 expose_value=False,
             )(f)
+        if enable_initial_state_element:
+            f = o(
+                "-e",
+                "--initial-state-element-override",
+                type=str,
+                multiple=True,
+                default=(),
+                help="Elements to override in the used initial state in the format DESIGN_FORMAT_ID=PATH",
+            )(f)
         if accept_config_files:
             f = argument(
                 "config_files",