Introduce shell state snapshots to autoenv plugin

This patch provides `_autoenv_snap_pre` & `_autoenv_snap_post` to create
a shell state diff for use in `.enter` files which can then be restored
in `.exit` with `_autoenv_snap_restore`.

autoenv/autoenv.zsh

@@ -132,6 +132,90 @@ commands:
 # Global entered directories array.
 _autoenv_entered=()
 
+# Per-directory shell state snapshots for computing enter/exit diffs.
+typeset -gA _autoenv_snap_pre_funcs _autoenv_snap_pre_aliases _autoenv_snap_pre_env
+typeset -gA _autoenv_snap_pre_path _autoenv_snap_pre_ps1 _autoenv_snap_pre_bindkeys
+typeset -gA _autoenv_snap_diff_funcs _autoenv_snap_diff_aliases _autoenv_snap_diff_env
+typeset -gA _autoenv_snap_new_bindkeys
+
+# Snapshot shell state before .enter modifications. Call in .enter BEFORE
+# sourcing scripts: _autoenv_snap_pre ${0:A:h}
+_autoenv_snap_pre() {
+  local dir=$1
+  _autoenv_snap_pre_funcs[$dir]="${(pj:\n:)${(@k)functions}}"
+  _autoenv_snap_pre_aliases[$dir]="${(pj:\n:)${(@k)aliases}}"
+  _autoenv_snap_pre_env[$dir]="${(F)$(typeset +gx 2>/dev/null)}"
+  _autoenv_snap_pre_path[$dir]="$PATH"
+  _autoenv_snap_pre_ps1[$dir]="$PS1"
+  _autoenv_snap_pre_bindkeys[$dir]="$(bindkey -L)"
+}
+
+# Snapshot shell state after .enter modifications and compute the diff.
+# Call in .enter AFTER sourcing scripts: _autoenv_snap_post ${0:A:h}
+_autoenv_snap_post() {
+  local dir=$1 f a v
+  local -a pre_funcs=("${(f)_autoenv_snap_pre_funcs[$dir]}")
+  local -a pre_aliases=("${(f)_autoenv_snap_pre_aliases[$dir]}")
+  local -a pre_env=("${(f)_autoenv_snap_pre_env[$dir]}")
+  local -a diff_funcs=() diff_aliases=() diff_env=()
+  for f in "${(@k)functions}"; do
+    (( ${pre_funcs[(Ie)$f]} )) || diff_funcs+=$f
+  done
+  for a in "${(@k)aliases}"; do
+    (( ${pre_aliases[(Ie)$a]} )) || diff_aliases+=$a
+  done
+  for v in "${(f)$(typeset +gx 2>/dev/null)}"; do
+    (( ${pre_env[(Ie)$v]} )) || diff_env+=$v
+  done
+  _autoenv_snap_diff_funcs[$dir]="${(pj:\n:)diff_funcs}"
+  _autoenv_snap_diff_aliases[$dir]="${(pj:\n:)diff_aliases}"
+  _autoenv_snap_diff_env[$dir]="${(pj:\n:)diff_env}"
+  # Diff bindkeys: find truly new bindings (key sequence didn't exist in pre).
+  # Changed bindings (same key, different widget) are restored via eval of pre.
+  local -a pre_bk=("${(f)_autoenv_snap_pre_bindkeys[$dir]}")
+  typeset -A pre_keyspecs
+  local bkline
+  for bkline in "${pre_bk[@]}"; do
+    [[ -n "$bkline" ]] && pre_keyspecs[${bkline% *}]=1
+  done
+  local -a new_bk=()
+  for bkline in "${(f)$(bindkey -L)}"; do
+    (( ${pre_bk[(Ie)$bkline]} )) && continue
+    (( ${+pre_keyspecs[${bkline% *}]} )) || new_bk+="$bkline"
+  done
+  _autoenv_snap_new_bindkeys[$dir]="${(pj:\n:)new_bk}"
+  unset "_autoenv_snap_pre_funcs[$dir]" "_autoenv_snap_pre_aliases[$dir]" \
+        "_autoenv_snap_pre_env[$dir]"
+}
+
+# Restore shell state by removing only what was added between snap_pre and
+# snap_post. Call in .exit: _autoenv_snap_restore ${0:A:h}
+_autoenv_snap_restore() {
+  local dir=$1 f a v
+  for f in "${(f)_autoenv_snap_diff_funcs[$dir]}"; do
+    [[ -n "$f" ]] && unset -f "$f"
+  done
+  for a in "${(f)_autoenv_snap_diff_aliases[$dir]}"; do
+    [[ -n "$a" ]] && unalias "$a" 2>/dev/null
+  done
+  for v in "${(f)_autoenv_snap_diff_env[$dir]}"; do
+    [[ -n "$v" ]] && unset "$v"
+  done
+  PATH="${_autoenv_snap_pre_path[$dir]}"
+  PS1="${_autoenv_snap_pre_ps1[$dir]}"
+  # Restore changed bindkeys to their pre values.
+  eval "${_autoenv_snap_pre_bindkeys[$dir]}"
+  # Remove truly new bindkeys (key sequence didn't exist before .enter).
+  local bkline
+  for bkline in "${(f)_autoenv_snap_new_bindkeys[$dir]}"; do
+    [[ -n "$bkline" ]] && eval "${${bkline% *}/bindkey/bindkey -r}"
+  done
+  unset "_autoenv_snap_pre_path[$dir]" "_autoenv_snap_pre_ps1[$dir]" \
+        "_autoenv_snap_pre_bindkeys[$dir]" "_autoenv_snap_new_bindkeys[$dir]" \
+        "_autoenv_snap_diff_funcs[$dir]" "_autoenv_snap_diff_aliases[$dir]" \
+        "_autoenv_snap_diff_env[$dir]"
+}
+
 # Load zstat from stat module for inspecting modified time.
 zmodload -F zsh/stat b:zstat
 

worktree/worktree.plugin.zsh

@@ -33,12 +33,16 @@ Manage git worktrees by branch name."
     add)
       local root=$(git rev-parse --show-toplevel)
       local wt_dest=${root:h}/${root:t}@${branch}
-      git worktree add "$wt_dest" "$branch"
+      if ! git show-ref --verify --quiet refs/heads/"$branch" && \
+         ! git show-ref --verify --quiet refs/remotes/origin/"$branch"; then
+        git branch "$branch" origin/main || return 1
+      fi
+      git worktree add "$wt_dest" "$branch" || return 1
       if [ -f $root/.enter ] && [ -f $root/.exit ]; then
         ln -s $root/.enter $wt_dest/.enter
         ln -s $root/.exit $wt_dest/.exit
-        _autoenv_authorized $PWD/.enter yes
+        _autoenv_authorized $wt_dest/.enter yes
-        _autoenv_authorized $PWD/.exit yes
+        _autoenv_authorized $wt_dest/.exit yes
       fi
       ;;
     rm|remove)