From eca7b960f684b9a226543612d5f08139061b380e Mon Sep 17 00:00:00 2001
From: abejgonzalez <abe.j.gonza@gmail.com>
Date: Mon, 6 Nov 2023 17:22:05 -0800
Subject: [PATCH] More robust env.sh | Fix conda activate if in subshell

---
 scripts/build-setup.sh                        | 22 ++++++--
 .../init-submodules-no-riscv-tools-nolog.sh   |  6 +-
 scripts/replace-content.py                    | 56 +++++++++++++++++++
 scripts/utils.sh                              | 24 ++++++++
 4 files changed, 100 insertions(+), 8 deletions(-)
 create mode 100755 scripts/replace-content.py

diff --git a/scripts/build-setup.sh b/scripts/build-setup.sh
index 3624b5e7ab..2c682514a1 100755
--- a/scripts/build-setup.sh
+++ b/scripts/build-setup.sh
@@ -250,11 +250,25 @@ if run_step "10"; then
     exit_if_last_command_failed
 fi
 
-cat <<EOT >> env.sh
-# line auto-generated by $0
+# Conda Setup
+# Provide a sourceable snippet that can be used in subshells that may not have
+# inhereted conda functions that would be brought in under a login shell that
+# has run conda init (e.g., VSCode, CI)
+read -r -d '\0' CONDA_ACTIVATE_PREAMBLE <<'END_CONDA_ACTIVATE'
+if ! type conda >& /dev/null; then
+    echo "::ERROR:: you must have conda in your environment first"
+    return 1  # don't want to exit here because this file is sourced
+fi
+
+# if we're sourcing this in a sub process that has conda in the PATH but not as a function, init it again
+conda activate --help >& /dev/null || source $(conda info --base)/etc/profile.d/conda.sh
+\0
+END_CONDA_ACTIVATE
+
+replace_content env.sh build-setup "# line auto-generated by $0
+$CONDA_ACTIVATE_PREAMBLE
 conda activate $CYDIR/.conda-env
-source $CYDIR/scripts/fix-open-files.sh
-EOT
+source $CYDIR/scripts/fix-open-files.sh"
 
 echo "Setup complete!"
 
diff --git a/scripts/init-submodules-no-riscv-tools-nolog.sh b/scripts/init-submodules-no-riscv-tools-nolog.sh
index d15498ea68..aa921f0d9a 100755
--- a/scripts/init-submodules-no-riscv-tools-nolog.sh
+++ b/scripts/init-submodules-no-riscv-tools-nolog.sh
@@ -157,8 +157,6 @@ if [ ! -f ./software/firemarshal/marshal-config.yaml ]; then
   echo "firesim-dir: '../../sims/firesim/'" > ./software/firemarshal/marshal-config.yaml
 fi
 
-cat << EOT >> env.sh
-# line auto-generated by init-submodules-no-riscv-tools.sh
+replace_content env.sh init-submodules "# line auto-generated by init-submodules-no-riscv-tools.sh
 __DIR="$RDIR"
-PATH=\$__DIR/software/firemarshal:\$PATH
-EOT
+PATH=\$__DIR/software/firemarshal:\$PATH"
diff --git a/scripts/replace-content.py b/scripts/replace-content.py
new file mode 100755
index 0000000000..4699ccca32
--- /dev/null
+++ b/scripts/replace-content.py
@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+
+# Replace text in a file given a key identifying a block to replace.
+# If the file doesn't exist, create it.
+#
+# args
+# $1 - file to replace text in
+# $2 - key used to find block of text to replace
+# $3 - text to fill in block that is replaced
+
+import re
+import sys
+
+def CY_INITIALIZE_RE_BLOCK(k):
+    return (
+        r"^# >>> " + f"{k}" + r" initialize >>>(?:\n|\r\n)"
+        r"([\s\S]*?)"
+        r"# <<< " + f"{k}" + r" initialize <<<(?:\n|\r\n)?"
+    )
+
+def CY_INITIALIZE_START_TOKEN(k):
+    return "# >>> " + f"{k}" + " initialize >>>"
+
+def CY_INITIALIZE_END_TOKEN(k):
+    return "# <<< " + f"{k}" + " initialize <<<"
+
+# ------------------------------
+
+try:
+    with open(sys.argv[1]) as fh:
+        fh_content = fh.read()
+except FileNotFoundError:
+    fh_content = ""
+except:
+    raise
+
+initialize_comment_key = sys.argv[2]
+inner_contents = CY_INITIALIZE_START_TOKEN(initialize_comment_key) + "\n" + sys.argv[3] + "\n" + CY_INITIALIZE_END_TOKEN(initialize_comment_key) + "\n"
+
+# ------------------------------
+
+replace_str = "__CY_REPLACE_ME_123__"
+fh_content = re.sub(
+    CY_INITIALIZE_RE_BLOCK(initialize_comment_key),
+    replace_str,
+    fh_content,
+    flags=re.MULTILINE,
+)
+# TODO: maybe remove all but last of replace_str, if there's more than one occurrence
+fh_content = fh_content.replace(replace_str, inner_contents)
+
+if CY_INITIALIZE_START_TOKEN(initialize_comment_key) not in fh_content:
+    fh_content += "\n%s\n" % inner_contents
+
+with open(sys.argv[1], "w") as fh:
+    fh.write(fh_content)
diff --git a/scripts/utils.sh b/scripts/utils.sh
index 30d3e092c2..633be1033e 100755
--- a/scripts/utils.sh
+++ b/scripts/utils.sh
@@ -53,3 +53,27 @@ function restore_bash_options
 {
     set +vx; eval "$OLDSTATE"
 }
+
+#######################################
+# Wrapper around replace-content.py.
+# For a file ($1), write out text ($3) into it
+# replacing any area designated by $2.
+#######################################
+function replace_content
+{
+    # On macOS, use GNU readlink from 'coreutils' package in Homebrew/MacPorts
+    if [ "$(uname -s)" = "Darwin" ] ; then
+        READLINK=greadlink
+    else
+        READLINK=readlink
+    fi
+
+    # If BASH_SOURCE is undefined, we may be running under zsh, in that case
+    # provide a zsh-compatible alternative
+    DIR="$(dirname "$($READLINK -f "${BASH_SOURCE[0]:-${(%):-%x}}")")"
+    file="$1"
+    shift
+    key="$1"
+    shift
+    $DIR/replace-content.py "$file" "$key" "$@"
+}