package cd.transform.analysis;

import cd.ir.Ast.*;
import cd.ir.AstVisitor;
import cd.ir.Symbol.ArrayTypeSymbol;
import cd.ir.Symbol.ClassSymbol;
import cd.ir.Symbol.TypeSymbol;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class ClassUsage extends AstVisitor<Void, Set<ClassSymbol>> {

	public Set<ClassDecl> go(List<ClassDecl> classDecls) {
		Set<ClassDecl> used = new HashSet<>();
		Set<ClassSymbol> todo = new HashSet<>();
		for (ClassDecl classDecl : classDecls)
			if (classDecl.name.equals("Main"))
				todo.add(classDecl.sym);

		while (!todo.isEmpty()) {
			ClassSymbol sym = todo.iterator().next();
			todo.remove(sym);
			if (!used.contains(sym.ast) && sym != ClassSymbol.objectType) {
				visit(sym.ast, todo);
				used.add(sym.ast);
			}
		}

		return used;
	}

	@Override
	public Void classDecl(ClassDecl ast, Set<ClassSymbol> arg) {
		arg.add(ast.sym.superClass);
		return super.classDecl(ast, arg);
	}

	@Override
	protected Void dfltExpr(Expr ast, Set<ClassSymbol> arg) {
		queueIfClassSymbol(ast.type, arg);
		return super.dfltExpr(ast, arg);
	}

	@Override
	public Void newObject(NewObject ast, Set<ClassSymbol> arg) {
		arg.add((ClassSymbol) ast.type);
		return super.newObject(ast, arg);
	}

	@Override
	public Void newArray(NewArray ast, Set<ClassSymbol> arg) {
		queueIfClassSymbol(ast.type, arg);
		return super.newArray(ast, arg);
	}

	@Override
	public Void varDecl(VarDecl ast, Set<ClassSymbol> arg) {
		queueIfClassSymbol(ast.sym.type, arg);
		return super.varDecl(ast, arg);
	}

	@Override
	public Void methodDecl(MethodDecl ast, Set<ClassSymbol> arg) {
		queueIfClassSymbol(ast.sym.returnType, arg);
		return super.methodDecl(ast, arg);
	}

	private void queueIfClassSymbol(TypeSymbol type, Set<ClassSymbol> arg) {
		if (type instanceof ArrayTypeSymbol)
			type = ((ArrayTypeSymbol) type).elementType;
		if (type instanceof ClassSymbol)
			arg.add((ClassSymbol) type);
	}
}
