// Copyright (c) .NET Foundation and contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
using System.Collections.Generic;
using System.Diagnostics;
using Mono.Cecil;
using Mono.Collections.Generic;
 
namespace Mono.Linker.Steps
{
	public class SealerStep : BaseStep
	{
		HashSet<TypeDefinition>? referencedBaseTypeCache;
 
		public SealerStep ()
		{
		}
 
		protected override void ProcessAssembly (AssemblyDefinition assembly)
		{
			if (!Context.CanApplyOptimization (CodeOptimizations.Sealer, assembly))
				return;
 
			foreach (var type in assembly.MainModule.Types)
				ProcessType (type);
		}
 
		protected override void EndProcess ()
		{
			referencedBaseTypeCache = null;
		}
 
		bool IsSubclassed (TypeDefinition type)
		{
			if (referencedBaseTypeCache == null) {
				referencedBaseTypeCache = new HashSet<TypeDefinition> ();
				foreach (var a in Context.GetAssemblies ()) {
					if (!Annotations.IsMarked (a))
						continue;
 
					PopulateCache (a.MainModule.Types);
 
					void PopulateCache (Collection<TypeDefinition> types)
					{
						foreach (var t in types) {
							var btd = Context.TryResolve (t.BaseType);
							if (btd != null)
								referencedBaseTypeCache.Add (btd);
 
							if (t.HasNestedTypes) {
								PopulateCache (t.NestedTypes);
							}
						}
					}
				}
			}
 
			var bt = Context.TryResolve (type);
			return bt != null && referencedBaseTypeCache.Contains (bt);
		}
 
		void ProcessType (TypeDefinition type)
		{
			if (type.HasNestedTypes) {
				foreach (var nt in type.NestedTypes) {
					ProcessType (nt);
				}
			}
 
			//
			// interface members are virtual (and we cannot change this)
			//
			if (type.IsInterface)
				return;
 
			//
			// the code does not include any subclass for this type
			//
			if (!type.IsAbstract && !type.IsSealed && !IsSubclassed (type))
				SealType (type);
 
			if (!type.HasMethods)
				return;
 
			// process methods to see if we can seal or devirtualize them
			foreach (var method in type.Methods) {
				if (method.IsFinal || !method.IsVirtual || method.IsAbstract || method.IsRuntime)
					continue;
 
				Debug.Assert (Annotations.IsMarked (method));
				if (!Annotations.IsMarked (method))
					continue;
 
				var overrides = Annotations.GetOverrides (method);
 
				//
				// cannot de-virtualize nor seal methods if something overrides them
				//
				if (IsAnyOverrideMarked (overrides))
					continue;
 
				SealMethod (method);
 
				// subclasses might need this method to satisfy an interface requirement
				// and requires dispatch/virtual support
				if (!type.IsSealed)
					continue;
 
				var bases = Annotations.GetBaseMethods (method);
				// Devirtualize if a method is not override to existing marked methods
				if (!IsAnyBaseMarked (bases))
					method.IsVirtual = method.IsFinal = method.IsNewSlot = false;
			}
		}
 
		protected virtual void SealType (TypeDefinition type)
		{
			type.IsSealed = true;
		}
 
		protected virtual void SealMethod (MethodDefinition method)
		{
			method.IsFinal = true;
		}
 
		bool IsAnyOverrideMarked (IEnumerable<OverrideInformation>? list)
		{
			if (list == null)
				return false;
 
			foreach (var m in list) {
				if (Annotations.IsMarked (m.Override))
					return true;
			}
			return false;
		}
 
		bool IsAnyBaseMarked (IEnumerable<OverrideInformation>? list)
		{
			if (list == null)
				return false;
 
			foreach (var m in list) {
				if (Annotations.IsMarked (m.Base))
					return true;
			}
			return false;
		}
	}
}