PowerShell · lzybkr · Nov 2, 2017 · Nov 1, 2017 · Nov 1, 2017 · daxian-dbw
diff --git a/src/System.Management.Automation/engine/Attributes.cs b/src/System.Management.Automation/engine/Attributes.cs
@@ -560,6 +560,10 @@ public string[] ParameterSetName
     [AttributeUsage(AttributeTargets.Assembly)]
     public class DynamicClassImplementationAssemblyAttribute : Attribute
     {
+        /// <summary>
+        /// The (possibly null) path to the file defining this class.
+        /// </summary>
+        public string ScriptFile { get; set; }
     }
 
     #endregion Misc Attributes

@@ -16,10 +16,11 @@ namespace System.Management.Automation.Language
 {
     internal class TypeDefiner
     {
+        internal const string DynamicClassAssemblyName = "PowerShell Class Assembly";
+
         private static int s_globalCounter = 0;
-        private static readonly object[] s_emptyArgArray = Utils.EmptyArray<object>();
         private static readonly CustomAttributeBuilder s_hiddenCustomAttributeBuilder =
-            new CustomAttributeBuilder(typeof(HiddenAttribute).GetConstructor(Type.EmptyTypes), s_emptyArgArray);
+            new CustomAttributeBuilder(typeof(HiddenAttribute).GetConstructor(Type.EmptyTypes), Utils.EmptyArray<object>());
 
         private static readonly string s_sessionStateKeeperFieldName = "__sessionStateKeeper";
         internal static readonly string SessionStateFieldName = "__sessionState";
@@ -1100,30 +1101,40 @@ internal void DefineEnum()
             }
         }
 
-        private static IEnumerable<CustomAttributeBuilder> GetAssemblyAttributeBuilders()
+        private static IEnumerable<CustomAttributeBuilder> GetAssemblyAttributeBuilders(string scriptFile)
         {
-            yield return new CustomAttributeBuilder(typeof(DynamicClassImplementationAssemblyAttribute).GetConstructor(Type.EmptyTypes), s_emptyArgArray);
+            var ctor = typeof(DynamicClassImplementationAssemblyAttribute).GetConstructor(Type.EmptyTypes);
+            var emptyArgs = Utils.EmptyArray<object>();
+
+            if (string.IsNullOrEmpty(scriptFile)) {
+                yield return new CustomAttributeBuilder(ctor, emptyArgs);
+                yield break;
+            }
+
+            var propertyInfo = new PropertyInfo[] {
+                typeof(DynamicClassImplementationAssemblyAttribute).GetProperty(nameof(DynamicClassImplementationAssemblyAttribute.ScriptFile)) };
+            var propertyArgs = new object[] { scriptFile };
+
+            yield return new CustomAttributeBuilder(ctor, emptyArgs,
+                propertyInfo, propertyArgs, Utils.EmptyArray<FieldInfo>(), emptyArgs);
+
         }
 
+        private static int counter = 0;
         internal static Assembly DefineTypes(Parser parser, Ast rootAst, TypeDefinitionAst[] typeDefinitions)
         {
             Diagnostics.Assert(rootAst.Parent == null, "Caller should only define types from the root ast");
 
             var definedTypes = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
 
-            // First character is a special mark that allows us to cheaply ignore dynamic generated assemblies in ClrFacade.GetAssemblies()
-            // The replaces at the end are for not-allowed characters. They are replaced by similar-looking chars.
-            string assemblyName = ClrFacade.FIRST_CHAR_PSASSEMBLY_MARK + (string.IsNullOrWhiteSpace(rootAst.Extent.File)
-                                      ? "powershell"
-                                      : rootAst.Extent.File
-                                      .Replace('\\', (char)0x29f9)
-                                      .Replace('/', (char)0x29f9)
-                                      .Replace(',', (char)0x201a)
-                                      .Replace(':', (char)0x0589));
-
-            var assembly = AssemblyBuilder.DefineDynamicAssembly(new AssemblyName(assemblyName),
-                                                                 AssemblyBuilderAccess.RunAndCollect, GetAssemblyAttributeBuilders());
-            var module = assembly.DefineDynamicModule(assemblyName);
+            var assemblyName = new AssemblyName(DynamicClassAssemblyName)
+            {
+                // We could generate a unique name, but a unique version works too.
+                Version = new Version(1, 0, 0, Interlocked.Increment(ref counter))
+            };
+            var assembly = AssemblyBuilder.DefineDynamicAssembly(assemblyName,
+                AssemblyBuilderAccess.RunAndCollect, GetAssemblyAttributeBuilders(rootAst.Extent.File));
+            var module = assembly.DefineDynamicModule(DynamicClassAssemblyName);
 
             var defineTypeHelpers = new List<DefineTypeHelper>();
             var defineEnumHelpers = new List<DefineEnumHelper>();

diff --git a/src/System.Management.Automation/utils/ClrFacade.cs b/src/System.Management.Automation/utils/ClrFacade.cs
@@ -41,12 +41,6 @@ static ClrFacade()
             }
         }
 
-        /// <summary>
-        /// We need it to avoid calling lookups inside dynamic assemblies with PS Types, so we exclude it from GetAssemblies().
-        /// We use this convention for names to archive it.
-        /// </summary>
-        internal static readonly char FIRST_CHAR_PSASSEMBLY_MARK = (char)0x29f9;
-
         #region Assembly
 
         internal static IEnumerable<Assembly> GetAssemblies(TypeResolutionState typeResolutionState, TypeName typeName)
@@ -65,7 +59,8 @@ internal static IEnumerable<Assembly> GetAssemblies(TypeResolutionState typeReso
         internal static IEnumerable<Assembly> GetAssemblies(string namespaceQualifiedTypeName = null)
         {
             return PSAssemblyLoadContext.GetAssembly(namespaceQualifiedTypeName) ??
-            AppDomain.CurrentDomain.GetAssemblies().Where(a => !(a.FullName.Length > 0 && a.FullName[0] == FIRST_CHAR_PSASSEMBLY_MARK));
+                   AppDomain.CurrentDomain.GetAssemblies().Where(a =>
+                       !TypeDefiner.DynamicClassAssemblyName.Equals(a.GetName().Name, StringComparison.Ordinal));
         }
 
         /// <summary>

diff --git a/test/powershell/Language/Classes/Scripting.Classes.BasicParsing.Tests.ps1 b/test/powershell/Language/Classes/Scripting.Classes.BasicParsing.Tests.ps1
@@ -688,6 +688,14 @@ Describe 'Type building' -Tags "CI" {
             ++$a::a | Should Be 2
         }
     }
+
+    It 'should get the script from a class type' {
+        class C {}
+
+        $a = [C].Assembly.GetCustomAttributes($false).Where{
+            $_ -is [System.Management.Automation.DynamicClassImplementationAssemblyAttribute]}
+        $a.ScriptFile | Should BeExactly $PSCommandPath
+    }
 }
 
 Describe 'RuntimeType created for TypeDefinitionAst' -Tags "CI" {