From 7cb9c08b754c280db7d1e03342b16e59d6214e61 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Jare=C5=A1?= <me@jakubjares.com>
Date: Thu, 23 May 2024 23:23:23 +0200
Subject: [PATCH] Merge changes from Profiler so we can share the code easily

---
 src/csharp/Pester/Pester.csproj               |   4 +
 .../Pester/Tracing/CodeCoverageTracer.cs      |   6 +-
 .../Pester/Tracing/ExternalTracerAdapter.cs   |  63 +--
 src/csharp/Pester/Tracing/ITracer.cs          |  13 +-
 src/csharp/Pester/Tracing/Tracer.cs           | 447 +++++++++---------
 src/csharp/Pester/Tracing/TracerHostUI.cs     | 173 +++----
 6 files changed, 352 insertions(+), 354 deletions(-)

diff --git a/src/csharp/Pester/Pester.csproj b/src/csharp/Pester/Pester.csproj
index 131cf867e..6392dbb9f 100644
--- a/src/csharp/Pester/Pester.csproj
+++ b/src/csharp/Pester/Pester.csproj
@@ -7,6 +7,10 @@
     <DebugType>embedded</DebugType>
   </PropertyGroup>
 
+  <PropertyGroup>
+    <DefineConstants>$(DefineConstants);PESTER</DefineConstants>
+  </PropertyGroup>
+
   <!-- PowerShell 7.2.x is the oldest supported PowerShell version. That version is built using net6.0.
   But there is a bug in 7.2.0 reference assemblies, up to 7.2.10, where the IExtens.File is missing from the reference assembly:
   https://github.com/PowerShell/PowerShell/issues/16408
diff --git a/src/csharp/Pester/Tracing/CodeCoverageTracer.cs b/src/csharp/Pester/Tracing/CodeCoverageTracer.cs
index 2a377bdaf..17d97da46 100644
--- a/src/csharp/Pester/Tracing/CodeCoverageTracer.cs
+++ b/src/csharp/Pester/Tracing/CodeCoverageTracer.cs
@@ -48,7 +48,7 @@ public CodeCoverageTracer(List<CodeCoveragePoint> points)
         // keyed as path -> line:column -> CodeCoveragePoint
         public Dictionary<string, Dictionary<string, List<CodeCoveragePoint>>> Hits { get; } = new Dictionary<string, Dictionary<string, List<CodeCoveragePoint>>>(StringComparer.OrdinalIgnoreCase);
 
-        public void Trace(string message, IScriptExtent extent, ScriptBlock _, int __)
+        public void Trace(string message, IScriptExtent extent, ScriptBlock _, int __, string ___, string ____)
         {
             if (_debug && extent?.File != null && CultureInfo.InvariantCulture.CompareInfo.IndexOf(extent.File, _debugFile, CompareOptions.OrdinalIgnoreCase) >= 0)
             {
@@ -106,10 +106,10 @@ public void Trace(string message, IScriptExtent extent, ScriptBlock _, int __)
 
 #pragma warning disable IDE0060
         // Profiler v3.1 compatible overload
-        public void Trace(IScriptExtent extent, ScriptBlock _, int __) => Trace(null, extent, _, __);
+        public void Trace(IScriptExtent extent, ScriptBlock scriptBlock, int level) => Trace(null, extent, scriptBlock, level, null, null);
 
         // Profiler v4 compatible overload
-        public void Trace(IScriptExtent extent, ScriptBlock _, int __, string ___, string ____) => Trace(null, extent, _, __);
+        public void Trace(IScriptExtent extent, ScriptBlock scriptBlock, int level, string functionName, string moduleName) => Trace(null, extent, scriptBlock, level, functionName, moduleName);
 #pragma warning restore IDE0060
     }
 }
diff --git a/src/csharp/Pester/Tracing/ExternalTracerAdapter.cs b/src/csharp/Pester/Tracing/ExternalTracerAdapter.cs
index 3cb057e91..87240a76d 100644
--- a/src/csharp/Pester/Tracing/ExternalTracerAdapter.cs
+++ b/src/csharp/Pester/Tracing/ExternalTracerAdapter.cs
@@ -3,49 +3,54 @@
 using System.Management.Automation.Language;
 using System.Reflection;
 
-namespace Pester.Tracing
+
+#if PESTER
+namespace Pester.Tracing;
+#else
+namespace Profiler;
+#endif
+
+class ExternalTracerAdapter : ITracer
 {
-    class ExternalTracerAdapter : ITracer
-    {
-        private readonly object _tracer;
-        private readonly MethodInfo _traceMethod;
-        private readonly int _version = 0;
+    private readonly object _tracer;
+    private readonly MethodInfo _traceMethod;
+    private readonly int _version = 0;
 
-        public object Tracer => _tracer;
+    public object Tracer => _tracer;
 
-        public ExternalTracerAdapter(object tracer)
-        {
-            // We got tracer that is not using the same types that we use here. Find a method based on the signature
-            // and use that. This enables tracers to register even when they don't take dependency on our types e.g. Pester CC tracer.
+    public ExternalTracerAdapter(object tracer)
+    {
+        // We got tracer that is not using the same types that we use here. Find a method based on the signature
+        // and use that. This enables tracers to register even when they don't take dependency on our types e.g. Pester CC tracer.
 
-            _tracer = tracer ?? new NullReferenceException(nameof(tracer));
-            _version = 2;
-            var traceMethod = tracer.GetType().GetMethod("Trace", new Type[] {
+        _tracer = tracer ?? new NullReferenceException(nameof(tracer));
+        _version = 2;
+        var traceMethod = tracer.GetType().GetMethod("Trace", new Type[] {
                 typeof(string), // message
                 typeof(IScriptExtent), // extent
                 typeof(ScriptBlock), // scriptblock
                 typeof(int) }); // level
 
-            if (traceMethod == null)
-            {
-                _version = 1;
-                traceMethod = tracer.GetType().GetMethod("Trace", new Type[] {
+        if (traceMethod == null)
+        {
+            _version = 1;
+            traceMethod = tracer.GetType().GetMethod("Trace", new Type[] {
                 typeof(string), // message
                 typeof(IScriptExtent), // extent
                 typeof(ScriptBlock), // scriptblock
                 typeof(int) }); // level
-            }
-
-            _traceMethod = traceMethod ??
-                throw new InvalidOperationException("The provided tracer does not have Trace method with this signature: Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level) or Trace(IScriptExtent extent, ScriptBlock scriptBlock, int level)");
         }
 
-        public void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level)
-        {
-            var parameters = _version == 2
-                ? new object[] { message, extent, scriptBlock, level }
-                : new object[] { extent, scriptBlock, level };
-            _traceMethod.Invoke(_tracer, parameters);
-        }
+        _traceMethod = traceMethod ??
+            throw new InvalidOperationException("The provided tracer does not have Trace method with this signature: Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level) or Trace(IScriptExtent extent, ScriptBlock scriptBlock, int level)");
+    }
+
+    public void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level, string functionName, string moduleName)
+    {
+        var parameters = _version == 2
+            ? new object[] { message, extent, scriptBlock, level }
+            : new object[] { extent, scriptBlock, level };
+        _traceMethod.Invoke(_tracer, parameters);
     }
 }
+
diff --git a/src/csharp/Pester/Tracing/ITracer.cs b/src/csharp/Pester/Tracing/ITracer.cs
index 3ac2d3916..4e820ab9b 100644
--- a/src/csharp/Pester/Tracing/ITracer.cs
+++ b/src/csharp/Pester/Tracing/ITracer.cs
@@ -1,10 +1,13 @@
 using System.Management.Automation;
 using System.Management.Automation.Language;
 
-namespace Pester.Tracing
+# if PESTER
+namespace Pester.Tracing;
+#else
+namespace Profiler;
+#endif
+
+public interface ITracer
 {
-    public interface ITracer
-    {
-        void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level);
-    }
+    void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level, string functionName, string moduleName);
 }
diff --git a/src/csharp/Pester/Tracing/Tracer.cs b/src/csharp/Pester/Tracing/Tracer.cs
index 3bdd653cc..3b9fff08b 100644
--- a/src/csharp/Pester/Tracing/Tracer.cs
+++ b/src/csharp/Pester/Tracing/Tracer.cs
@@ -3,266 +3,249 @@
 using System.Management.Automation.Host;
 using System.Management.Automation.Language;
 using System.Reflection;
-using NonGeneric = System.Collections;
 
-namespace Pester.Tracing
+# if PESTER
+namespace Pester.Tracing;
+#else
+namespace Profiler;
+#endif
+
+public static class Tracer
 {
-    public static class Tracer
+    private static Func<TraceLineInfo> GetTraceLineInfo;
+    private static TraceLineInfo LastTraceItem;
+
+    private static Action ResetUI = () => { };
+    /// <summary>
+    /// Primary tracer (slot1).
+    /// </summary>
+    public static ITracer Tracer1 { get; private set; }
+
+    /// <summary>
+    /// Secondary tracer (slot2).
+    /// </summary>
+    public static ITracer Tracer2 { get; private set; }
+
+    /// <summary>
+    /// Check if we should call Register or Patch functions for the given tracer. It returns false when tracer is not present in slot 1,
+    /// or if overwrite is true (default) and slot 1 has the same tracer as what we are registering. Use overwrite to check the type of the registered tracer
+    /// and allow replacing it with the new tracer if the types are the same. (e.g you are adding tracer for Pester code coverage, and user aborted the run
+    /// so Pester tracer from the previous run is still present, and you are now replacing it with another Pester code coverage tracer). You would rarely want
+    /// overwrite set to false. But for example when you run Profiler in Profiler, you would want to register the second tracer into slot2 not even though slot1
+    /// has tracer of the same type.
+    /// </summary>
+    /// <param name="tracer">The tracer to use.</param>
+    /// <param name="overwrite">Allow overwriting the tracer in the main tracer slot if the tracer type we are adding is the same as the one that is already present.</param>
+    /// <returns></returns>
+    public static bool ShouldRegisterTracer(object tracer, bool overwrite = true)
     {
-        private static Func<TraceLineInfo> GetTraceLineInfo;
-        private static Action ResetUI = () => { };
-        /// <summary>
-        /// Primary tracer (slot1).
-        /// </summary>
-        public static ITracer Tracer1 { get; private set; }
-        /// <summary>
-        /// Secondary tracer (slot2).
-        /// </summary>
-        public static ITracer Tracer2 { get; private set; }
-
-        [Obsolete("IsEnabled is obsolete because the internal state can become corrupted when script is cancelled, use ShouldRegisterTracer instead. The state of IsEnabled is set as before, but it won't be checked when using Patch or Register, to prevent the user session from being locked up by the incorrect state.")]
-        public static bool IsEnabled { get; private set; }
-
-        /// <summary>
-        /// Check if we should call Register or Patch functions for the given tracer. It returns false when tracer is not present in slot 1,
-        /// or if overwrite is true (default) and slot 1 has the same tracer as what we are registering. Use overwrite to check the type of the registered tracer
-        /// and allow replacing it with the new tracer if the types are the same. (e.g you are adding tracer for Pester code coverage, and user aborted the run
-        /// so Pester tracer from the previous run is still present, and you are now replacing it with another Pester code coverage tracer). You would rarely want
-        /// overwrite set to false. But for example when you run Profiler in Profiler, you would want to register the second tracer into slot2 not even though slot1
-        /// has tracer of the same type.
-        /// </summary>
-        /// <param name="tracer">The tracer to use.</param>
-        /// <param name="overwrite">Allow overwriting the tracer in the main tracer slot if the tracer type we are adding is the same as the one that is already present.</param>
-        /// <returns></returns>
-        public static bool ShouldRegisterTracer(object tracer, bool overwrite = true)
+        if (tracer is null)
+            throw new ArgumentNullException(nameof(tracer));
+
+        if (overwrite)
         {
-            if (tracer is null)
-                throw new ArgumentNullException(nameof(tracer));
-
-            if (overwrite)
-            {
-                // if there is a tracer in slot 1, and it is not the same as we are providing
-                // we should use slot 2 and only register the tracer
-                return HasDifferentTracer(Tracer1, tracer);
-            }
-
-            // we have tracer in slot 1, use slot 2 by registering the tracer
-            return HasTracer(Tracer1);
+            // if there is a tracer in slot 1, and it is not the same as we are providing
+            // we should use slot 2 and only register the tracer
+            return HasDifferentTracer(Tracer1, tracer);
         }
 
-        private static bool HasTracer(ITracer currentTracer)
+        // we have tracer in slot 1, use slot 2 by registering the tracer
+        return HasTracer(Tracer1);
+    }
+
+    private static bool HasTracer(ITracer currentTracer)
+    {
+        return currentTracer != null;
+    }
+
+    private static bool HasDifferentTracer(ITracer currentTracer, object newTracer)
+    {
+        if (currentTracer == null)
+            return false;
+
+        if (currentTracer is ExternalTracerAdapter adapted)
         {
-            return currentTracer != null;
+            return adapted.Tracer.GetType().Name != newTracer.GetType().Name;
         }
 
-        private static bool HasDifferentTracer(ITracer currentTracer, object newTracer)
+        return currentTracer.GetType().Name != newTracer.GetType().Name;
+    }
+
+    private static void EnsureTracingWasNotCorrupted()
+    {
+        // Code like Set-PSDebug -Trace 0, or Get-PSBreakPoint | Remove-PSBreakPoint will break the tracing, and we will lose data,
+        // the only line that is allowed to do that is $corruptionAutodetectionVariable = Set-PsDebug -Trace 0
+        // This line should be the second to last line that was executed. We re-enable the trace in Profiler, so the last line (::Unregister or ::Unpatch) is
+        // captured even when tracing is broken by user.
+        var corrupted = LastTraceItem.Extent?.Text != null && LastTraceItem.Extent.Text != "$corruptionAutodetectionVariable = Set-PsDebug -Trace 0";
+        if (corrupted)
         {
-            if (currentTracer == null)
-                return false;
+            throw new InvalidOperationException($"Trace was broken by: {LastTraceItem.Extent.Text} from {LastTraceItem.Extent.File}:{LastTraceItem.Extent.StartLineNumber}");
+        }
+    }
 
-            if (currentTracer is ExternalTracerAdapter adapted)
-            {
-                return adapted.Tracer.GetType().Name != newTracer.GetType().Name;
-            }
+    /// <summary>
+    /// Add a tracer to already setup session (to slot2). A tracer must implement ITracer or have a method `void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level)`.
+    /// </summary>
+    /// <param name="tracer"></param>
+    public static void Register(object tracer)
+    {
+        if (tracer is null)
+            throw new ArgumentNullException(nameof(tracer));
 
-            return currentTracer.GetType().Name != newTracer.GetType().Name;
-        }
+        if (!HasTracer(Tracer1))
+            throw new InvalidOperationException($"Tracer1 is null. If you want to activate tracing call {nameof(Patch)}.");
 
-        /// <summary>
-        /// Add a tracer to already setup session (to slot2). A tracer must implement ITracer or have a method `void Trace(string message, IScriptExtent extent, ScriptBlock scriptBlock, int level)`.
-        /// </summary>
-        /// <param name="tracer"></param>
-        public static void Register(object tracer)
-        {
-            if (tracer is null)
-                throw new ArgumentNullException(nameof(tracer));
+        if (HasDifferentTracer(Tracer2, tracer))
+            throw new InvalidOperationException($"Tracer2 already has tracer {Tracer2.GetType().Name}, and you are registering {tracer.GetType().Name}.");
 
-            if (!HasTracer(Tracer1))
-                throw new InvalidOperationException($"Tracer1 is null. If you want to activate tracing call {nameof(Patch)}.");
+        Tracer2 = tracer is ITracer t ? t : new ExternalTracerAdapter(tracer);
 
-            if (HasDifferentTracer(Tracer2, tracer))
-                throw new InvalidOperationException($"Tracer2 already has tracer {Tracer2.GetType().Name}, and you are registering {tracer.GetType().Name}.");
+        TraceLine(justTracer2: true);
+    }
 
-            Tracer2 = tracer is ITracer t ? t : new ExternalTracerAdapter(tracer);
+    /// <summary>
+    /// Unregister tracer from an already setup session (slot2).
+    /// </summary>
+    public static void Unregister()
+    {
+        TraceLine(justTracer2: true, storeTraceItem: false);
+        TraceLine(justTracer2: true, storeTraceItem: false);
+        Tracer2 = null;
 
-            TraceLine(justTracer2: true);
-        }
+        EnsureTracingWasNotCorrupted();
+    }
 
-        /// <summary>
-        /// Unregister tracer from an already setup session (slot2).
-        /// </summary>
-        public static void Unregister()
-        {
-            TraceLine(justTracer2: true);
-            TraceLine(justTracer2: true);
-            Tracer2 = null;
-        }
+    /// <summary>
+    /// Enable tracing by patch the current session by replacing the UI host with another that triggers tracing on every statement. Set-PSDebug -Trace 1 needs to be called before this, and Set-PSDebug -Off needs to be called after Unpatch.
+    /// </summary>
+    /// <param name="powerShellVersion">Major PSVersion that is used `$PSVersionTable.PSVersion.Major`</param>
+    /// <param name="context">ExecutionContext to be used `$ExecutionContext`</param>
+    /// <param name="ui">UIHost to be replaced. `$host.UI`</param>
+    /// <param name="tracer">The tracer to be used. For example ProfilerTracer.</param>
+    public static void Patch(int powerShellVersion, EngineIntrinsics context, PSHostUserInterface ui, ITracer tracer)
+    {
+        if (context is null)
+            throw new ArgumentNullException(nameof(context));
 
-        /// <summary>
-        /// Enable tracing by patch the current session by replacing the UI host with another that triggers tracing on every statement. Set-PSDebug -Trace 1 needs to be called before this, and Set-PSDebug -Off needs to be called after Unpatch.
-        /// </summary>
-        /// <param name="powerShellVersion">Major PSVersion that is used `$PSVersionTable.PSVersion.Major`</param>
-        /// <param name="context">ExecutionContext to be used `$ExecutionContext`</param>
-        /// <param name="ui">UIHost to be replaced. `$host.UI`</param>
-        /// <param name="tracer">The tracer to be used. For example ProfilerTracer.</param>
-        public static void Patch(int powerShellVersion, EngineIntrinsics context, PSHostUserInterface ui, ITracer tracer)
+        if (ui is null)
+            throw new ArgumentNullException(nameof(ui));
+
+        Tracer1 = tracer ?? throw new ArgumentNullException(nameof(tracer));
+
+        var uiFieldName = powerShellVersion >= 6 ? "_externalUI" : "externalUI";
+        // we get InternalHostUserInterface, grab external ui from that and replace it with ours
+        var externalUIField = ui.GetType().GetField(uiFieldName, BindingFlags.Instance | BindingFlags.NonPublic);
+        var externalUI = (PSHostUserInterface)externalUIField.GetValue(ui);
+
+        // replace it with out patched up UI that writes to profiler on debug
+        externalUIField.SetValue(ui, new TracerHostUI(externalUI, (message) => TraceLine(message, false)));
+
+        ResetUI = () => externalUIField.SetValue(ui, externalUI);
+
+        // getting MethodInfo of context._context.Debugger.TraceLine
+        var bf = BindingFlags.NonPublic | BindingFlags.Instance;
+        var contextInternal = context.GetType().GetField("_context", bf).GetValue(context);
+        var debugger = contextInternal.GetType().GetProperty("Debugger", bf).GetValue(contextInternal);
+        var debuggerType = debugger.GetType();
+
+        var callStackField = debuggerType.GetField("_callStack", BindingFlags.Instance | BindingFlags.NonPublic);
+        var _callStack = callStackField.GetValue(debugger);
+
+        var callStackType = _callStack.GetType();
+
+        var countBindingFlags = BindingFlags.Instance | BindingFlags.NonPublic;
+        var countProperty = callStackType.GetProperty("Count", countBindingFlags);
+        var getCount = countProperty.GetMethod;
+        var empty = new object[0];
+        var stack = callStackField.GetValue(debugger);
+        var initialLevel = (int)getCount.Invoke(stack, empty);
+
+        var lastFunctionContextMethod = callStackType.GetMethod("LastFunctionContext", BindingFlags.Instance | BindingFlags.NonPublic);
+
+        object functionContext1 = lastFunctionContextMethod.Invoke(callStackField.GetValue(debugger), empty);
+        var functionContextType = functionContext1.GetType();
+        var functionNameField = functionContextType.GetField("_functionName", BindingFlags.Instance | BindingFlags.NonPublic);
+        var executionContextField = functionContextType.GetField("_executionContext", BindingFlags.Instance | BindingFlags.NonPublic);
+        var sessionStateProperty = executionContextField.FieldType.GetProperty("SessionState", BindingFlags.Instance | BindingFlags.NonPublic);
+        var scriptBlockField = functionContextType.GetField("_scriptBlock", BindingFlags.Instance | BindingFlags.NonPublic);
+        var currentPositionProperty = functionContextType.GetProperty("CurrentPosition", BindingFlags.Instance | BindingFlags.NonPublic);
+
+        var scriptBlock1 = (ScriptBlock)scriptBlockField.GetValue(functionContext1);
+        var extent1 = (IScriptExtent)currentPositionProperty.GetValue(functionContext1);
+
+        GetTraceLineInfo = () =>
         {
-            if (context is null)
-                throw new ArgumentNullException(nameof(context));
-
-            if (ui is null)
-                throw new ArgumentNullException(nameof(ui));
-
-            Tracer1 = tracer ?? throw new ArgumentNullException(nameof(tracer));
-
-            var uiFieldName = powerShellVersion >= 6 ? "_externalUI" : "externalUI";
-            // we get InternalHostUserInterface, grab external ui from that and replace it with ours
-            var externalUIField = ui.GetType().GetField(uiFieldName, BindingFlags.Instance | BindingFlags.NonPublic);
-            var externalUI = (PSHostUserInterface)externalUIField.GetValue(ui);
-
-            // replace it with out patched up UI that writes to profiler on debug
-            externalUIField.SetValue(ui, new TracerHostUI(externalUI, (message) => TraceLine(message, false)));
-
-            ResetUI = () => externalUIField.SetValue(ui, externalUI);
-
-            // getting MethodInfo of context._context.Debugger.TraceLine
-            var bf = BindingFlags.NonPublic | BindingFlags.Instance;
-            var contextInternal = context.GetType().GetField("_context", bf).GetValue(context);
-            var debugger = contextInternal.GetType().GetProperty("Debugger", bf).GetValue(contextInternal);
-            var debuggerType = debugger.GetType();
-
-            var callStackField = debuggerType.GetField("_callStack", BindingFlags.Instance | BindingFlags.NonPublic);
-            var _callStack = callStackField.GetValue(debugger);
-
-            var callStackType = _callStack.GetType();
-
-            var countBindingFlags = BindingFlags.Instance | BindingFlags.NonPublic;
-            if (powerShellVersion == 3)
-            {
-                // in PowerShell 3 callstack is List<CallStackInfo> not a struct CallStackList
-                // Count is public property
-                countBindingFlags = BindingFlags.Instance | BindingFlags.Public;
-            }
-            var countProperty = callStackType.GetProperty("Count", countBindingFlags);
-            var getCount = countProperty.GetMethod;
-            var empty = new object[0];
-            var stack = callStackField.GetValue(debugger);
-            var initialLevel = (int)getCount.Invoke(stack, empty);
-
-            if (powerShellVersion == 3)
-            {
-                // we do the same operation as in the TraceLineAction below, but here
-                // we resolve the static things like types and properties, and then in the
-                // action we just use them to get the live data without the overhead of looking
-                // up properties all the time. This might be internally done in the reflection code
-                // did not measure the impact, and it is probably done for us in the reflection api itself
-                // in modern versions of runtime
-                var callStack1 = callStackField.GetValue(debugger);
-                var callStackList1 = (NonGeneric.IList)callStack1;
-                var level1 = callStackList1.Count - initialLevel;
-                var last1 = callStackList1[callStackList1.Count - 1];
-                var lastType = last1.GetType();
-                var functionContextProperty = lastType.GetProperty("FunctionContext", BindingFlags.NonPublic | BindingFlags.Instance);
-                var functionContext1 = functionContextProperty.GetValue(last1);
-                var functionContextType = functionContext1.GetType();
-
-                var scriptBlockField = functionContextType.GetField("_scriptBlock", BindingFlags.Instance | BindingFlags.NonPublic);
-                var currentPositionProperty = functionContextType.GetProperty("CurrentPosition", BindingFlags.Instance | BindingFlags.NonPublic);
-
-                var scriptBlock1 = (ScriptBlock)scriptBlockField.GetValue(functionContext1);
-                var extent1 = (IScriptExtent)currentPositionProperty.GetValue(functionContext1);
-
-                GetTraceLineInfo = () =>
-                {
-                    var callStack = callStackField.GetValue(debugger);
-                    var callStackList = (NonGeneric.IList)callStack;
-                    var level = callStackList.Count - initialLevel;
-                    var last = callStackList[callStackList.Count - 1];
-                    var functionContext = functionContextProperty.GetValue(last);
-
-                    var scriptBlock = (ScriptBlock)scriptBlockField.GetValue(functionContext);
-                    var extent = (IScriptExtent)currentPositionProperty.GetValue(functionContext);
-
-                    return new TraceLineInfo(extent, scriptBlock, level);
-                };
-            }
-            else
-            {
-                var lastFunctionContextMethod = callStackType.GetMethod("LastFunctionContext", BindingFlags.Instance | BindingFlags.NonPublic);
-
-                object functionContext1 = lastFunctionContextMethod.Invoke(callStackField.GetValue(debugger), empty);
-                var functionContextType = functionContext1.GetType();
-                var scriptBlockField = functionContextType.GetField("_scriptBlock", BindingFlags.Instance | BindingFlags.NonPublic);
-                var currentPositionProperty = functionContextType.GetProperty("CurrentPosition", BindingFlags.Instance | BindingFlags.NonPublic);
-
-                var scriptBlock1 = (ScriptBlock)scriptBlockField.GetValue(functionContext1);
-                var extent1 = (IScriptExtent)currentPositionProperty.GetValue(functionContext1);
-
-                GetTraceLineInfo = () =>
-                {
-                    var callStack = callStackField.GetValue(debugger);
-                    var level = (int)getCount.Invoke(callStack, empty) - initialLevel;
-                    object functionContext = lastFunctionContextMethod.Invoke(callStack, empty);
-                    var scriptBlock = (ScriptBlock)scriptBlockField.GetValue(functionContext);
-                    var extent = (IScriptExtent)currentPositionProperty.GetValue(functionContext);
-
-                    return new TraceLineInfo(extent, scriptBlock, level);
-                };
-            }
-
-#pragma warning disable CS0618 // Type or member is obsolete
-            IsEnabled = true;
-#pragma warning restore CS0618 // Type or member is obsolete
-
-            // Add another event to the top apart from the scriptblock invocation
-            // in Trace-ScriptInternal, this makes it more consistently work on first
-            // run. Without this, the triggering line sometimes does not show up as 99.9%
-            TraceLine();
-        }
+            var callStack = callStackField.GetValue(debugger);
+            var level = (int)getCount.Invoke(callStack, empty) - initialLevel;
+            object functionContext = lastFunctionContextMethod.Invoke(callStack, empty);
+            var executionContext = executionContextField.GetValue(functionContext);
+            var sessionState = (SessionState)sessionStateProperty.GetValue(executionContext, empty);
+            var moduleName = sessionState.Module?.Name;
+            var functionName = (string)functionNameField.GetValue(functionContext);
+            var scriptBlock = (ScriptBlock)scriptBlockField.GetValue(functionContext);
+            var extent = (IScriptExtent)currentPositionProperty.GetValue(functionContext);
+
+            return new TraceLineInfo(extent, scriptBlock, level, functionName, moduleName);
+        };
+
+        // Add another event to the top apart from the ScriptBlock invocation
+        // in Trace-ScriptInternal, this makes it more consistently work on first
+        // run. Without this, the triggering line sometimes does not show up as 99.9%
+        TraceLine();
+    }
+
+    /// <summary>
+    /// Put the original UI host in place and stop tracing.  Set-PSDebug -Trace 0 (or Set-PSDebug -Off) needs to be called before this.
+    /// </summary>
+    public static void Unpatch()
+    {
+        // Add Set-PSDebug -Trace 0 event and also another one for the internal disable
+        // this make first run more consistent for some reason
+        TraceLine(storeTraceItem: false);
+        TraceLine(storeTraceItem: false);
+        ResetUI();
+        Tracer1 = null;
+        Tracer2 = null;
+
+        EnsureTracingWasNotCorrupted();
+    }
 
-        /// <summary>
-        /// Put the original UI host in place and stop tracing.  Set-PSDebug -Trace 0 (or Set-PSDebug -Off) needs to be called before this.
-        /// </summary>
-        public static void Unpatch()
+    private static void TraceLine(string message = null, bool justTracer2 = false, bool storeTraceItem = true)
+    {
+        if (GetTraceLineInfo == null)
+            return;
+
+        var traceLineInfo = GetTraceLineInfo();
+
+        // Keep last tracing in memory so we can compare them when we unregister, or unpatch to detect corruption.
+        if (storeTraceItem)
         {
-#pragma warning disable CS0618 // Type or member is obsolete
-            IsEnabled = false;
-#pragma warning restore CS0618 // Type or member is obsolete
-
-            // Add Set-PSDebug -Trace 0 event and also another one for the internal disable
-            // this make first run more consistent for some reason
-            TraceLine();
-            TraceLine();
-            ResetUI();
-            Tracer1 = null;
-            Tracer2 = null;
+            LastTraceItem = traceLineInfo;
         }
 
-        private static void TraceLine(string message = null, bool justTracer2 = false)
+        if (!justTracer2)
         {
-            if (GetTraceLineInfo == null)
-                return;
-
-            var traceLineInfo = GetTraceLineInfo();
-            if (!justTracer2)
-            {
-                Tracer1?.Trace(message, traceLineInfo.Extent, traceLineInfo.ScriptBlock, traceLineInfo.Level);
-            }
-            Tracer2?.Trace(message, traceLineInfo.Extent, traceLineInfo.ScriptBlock, traceLineInfo.Level);
+            Tracer1?.Trace(message, traceLineInfo.Extent, traceLineInfo.ScriptBlock, traceLineInfo.Level, traceLineInfo.FunctionName, traceLineInfo.ModuleName);
         }
+        Tracer2?.Trace(message, traceLineInfo.Extent, traceLineInfo.ScriptBlock, traceLineInfo.Level, traceLineInfo.FunctionName, traceLineInfo.ModuleName);
+    }
+
+    private struct TraceLineInfo
+    {
+        public IScriptExtent Extent;
+        public ScriptBlock ScriptBlock;
+        public int Level;
+        public string FunctionName;
+        public string ModuleName;
 
-        private struct TraceLineInfo
+        public TraceLineInfo(IScriptExtent extent, ScriptBlock scriptBlock, int level, string functionName, string moduleName)
         {
-            public IScriptExtent Extent;
-            public ScriptBlock ScriptBlock;
-            public int Level;
-
-            public TraceLineInfo(IScriptExtent extent, ScriptBlock scriptBlock, int level)
-            {
-                Extent = extent;
-                ScriptBlock = scriptBlock;
-                Level = level;
-            }
+            Extent = extent;
+            ScriptBlock = scriptBlock;
+            Level = level;
+            FunctionName = functionName;
+            ModuleName = moduleName;
         }
     }
 }
diff --git a/src/csharp/Pester/Tracing/TracerHostUI.cs b/src/csharp/Pester/Tracing/TracerHostUI.cs
index d9f2afd18..9433c367a 100644
--- a/src/csharp/Pester/Tracing/TracerHostUI.cs
+++ b/src/csharp/Pester/Tracing/TracerHostUI.cs
@@ -5,92 +5,95 @@
 using System.Management.Automation.Host;
 using System.Security;
 
-namespace Pester.Tracing
+#if PESTER
+namespace Pester.Tracing;
+#else
+namespace Profiler;
+#endif
+
+internal class TracerHostUI : PSHostUserInterface
 {
-    internal class TracerHostUI : PSHostUserInterface
+    private readonly PSHostUserInterface _ui;
+    private readonly Action<string> _trace;
+
+    public TracerHostUI(PSHostUserInterface ui, Action<string> trace)
+    {
+        _ui = ui;
+        _trace = trace;
+    }
+
+    public override PSHostRawUserInterface RawUI => _ui.RawUI;
+
+    public override Dictionary<string, PSObject> Prompt(string caption, string message, Collection<FieldDescription> descriptions)
+    {
+        return _ui.Prompt(caption, message, descriptions);
+    }
+
+    public override int PromptForChoice(string caption, string message, Collection<ChoiceDescription> choices, int defaultChoice)
+    {
+        return _ui.PromptForChoice(caption, message, choices, defaultChoice);
+    }
+
+    public override PSCredential PromptForCredential(string caption, string message, string userName, string targetName)
+    {
+        return _ui.PromptForCredential(caption, message, userName, targetName);
+    }
+
+    public override PSCredential PromptForCredential(string caption, string message, string userName, string targetName, PSCredentialTypes allowedCredentialTypes, PSCredentialUIOptions options)
+    {
+        return _ui.PromptForCredential(caption, message, userName, targetName, allowedCredentialTypes, options);
+    }
+
+    public override string ReadLine()
+    {
+        return _ui.ReadLine();
+    }
+
+    public override SecureString ReadLineAsSecureString()
+    {
+        return _ui.ReadLineAsSecureString();
+    }
+
+    public override void Write(string value)
+    {
+        _ui.Write(value);
+    }
+
+    public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
+    {
+        _ui.Write(foregroundColor, backgroundColor, value);
+    }
+
+    public override void WriteDebugLine(string message)
+    {
+        if (_trace == null)
+            _ui.WriteDebugLine(message);
+
+        _trace(message);
+    }
+
+    public override void WriteErrorLine(string value)
+    {
+        _ui.WriteErrorLine(value);
+    }
+
+    public override void WriteLine(string value)
+    {
+        _ui.WriteLine(value);
+    }
+
+    public override void WriteProgress(long sourceId, ProgressRecord record)
+    {
+        _ui.WriteProgress(sourceId, record);
+    }
+
+    public override void WriteVerboseLine(string message)
+    {
+        _ui.WriteVerboseLine(message);
+    }
+
+    public override void WriteWarningLine(string message)
     {
-        private readonly PSHostUserInterface _ui;
-        private readonly Action<string> _trace;
-
-        public TracerHostUI(PSHostUserInterface ui, Action<string> trace)
-        {
-            _ui = ui;
-            _trace = trace;
-        }
-
-        public override PSHostRawUserInterface RawUI => _ui.RawUI;
-
-        public override Dictionary<string, PSObject> Prompt(string caption, string message, Collection<FieldDescription> descriptions)
-        {
-            return _ui.Prompt(caption, message, descriptions);
-        }
-
-        public override int PromptForChoice(string caption, string message, Collection<ChoiceDescription> choices, int defaultChoice)
-        {
-            return _ui.PromptForChoice(caption, message, choices, defaultChoice);
-        }
-
-        public override PSCredential PromptForCredential(string caption, string message, string userName, string targetName)
-        {
-            return _ui.PromptForCredential(caption, message, userName, targetName);
-        }
-
-        public override PSCredential PromptForCredential(string caption, string message, string userName, string targetName, PSCredentialTypes allowedCredentialTypes, PSCredentialUIOptions options)
-        {
-            return _ui.PromptForCredential(caption, message, userName, targetName, allowedCredentialTypes, options);
-        }
-
-        public override string ReadLine()
-        {
-            return _ui.ReadLine();
-        }
-
-        public override SecureString ReadLineAsSecureString()
-        {
-            return _ui.ReadLineAsSecureString();
-        }
-
-        public override void Write(string value)
-        {
-            _ui.Write(value);
-        }
-
-        public override void Write(ConsoleColor foregroundColor, ConsoleColor backgroundColor, string value)
-        {
-            _ui.Write(foregroundColor, backgroundColor, value);
-        }
-
-        public override void WriteDebugLine(string message)
-        {
-            if (_trace == null)
-                _ui.WriteDebugLine(message);
-
-            _trace(message);
-        }
-
-        public override void WriteErrorLine(string value)
-        {
-            _ui.WriteErrorLine(value);
-        }
-
-        public override void WriteLine(string value)
-        {
-            _ui.WriteLine(value);
-        }
-
-        public override void WriteProgress(long sourceId, ProgressRecord record)
-        {
-            _ui.WriteProgress(sourceId, record);
-        }
-
-        public override void WriteVerboseLine(string message)
-        {
-            _ui.WriteVerboseLine(message);
-        }
-
-        public override void WriteWarningLine(string message)
-        {
-            _ui.WriteWarningLine(message);
-        }
+        _ui.WriteWarningLine(message);
     }
 }