<commit_before>package de.peeeq.wurstscript.intermediateLang.optimizer;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.eclipse.jdt.annotation.NonNull;
import org.eclipse.jdt.annotation.Nullable;

import de.peeeq.wurstscript.intermediateLang.optimizer.ControlFlowGraph.Node;
import de.peeeq.wurstscript.jassIm.ImAlloc;
import de.peeeq.wurstscript.jassIm.ImConst;
import de.peeeq.wurstscript.jassIm.ImFunction;
import de.peeeq.wurstscript.jassIm.ImProg;
import de.peeeq.wurstscript.jassIm.ImSet;
import de.peeeq.wurstscript.jassIm.ImStmt;
import de.peeeq.wurstscript.jassIm.ImVar;
import de.peeeq.wurstscript.jassIm.ImVarAccess;
import de.peeeq.wurstscript.translation.imtranslation.ImTranslator;

public class ConstantAndCopyPropagation {

	private final ImProg prog;
	private final ImTranslator trans;

	public ConstantAndCopyPropagation(ImTranslator trans) {
		this.prog = trans.getImProg();
		this.trans = trans;
	}

	public void optimize() {
		for (ImFunction func : prog.getFunctions()) {
			optimizeFunc(func);
		}
	}

	static class Value {
		// one of the two is null
		final @Nullable ImVar copyVar;
		final @Nullable ImConst constantValue;

		public Value(ImVar copyVar) {
			this.copyVar = copyVar;
			this.constantValue = null;
		}

		public Value(ImConst constantValue) {
			this.copyVar = null;
			this.constantValue = constantValue;
		}
		
		@Override
		public boolean equals(@Nullable Object obj) {
			if (obj instanceof Value) {
				return equalValue((Value) obj);
			}
			return false;
		}

		public boolean equalValue(Value other) {
			if (copyVar == other.copyVar) {
				return true;
			} else if (constantValue != null && other.constantValue != null) {
				return constantValue.equalValue(other.constantValue);
			}
			return false;
		}
		
		@Override
		public String toString() {
			if (copyVar != null) {
				return "copy of " + copyVar;
			} else {
				return "constant " + constantValue;
			}
		}

	}

	static class Knowledge {

		Map<ImVar, Value> varKnowledge = new HashMap<>();
		Map<ImVar, Value> varKnowledgeOut = new HashMap<>();
		@Override
		public String toString() {
			return "[in =" + varKnowledge + ", out=" + varKnowledgeOut + "]";
		}
		
		

	}

	private void optimizeFunc(ImFunction func) {
		ControlFlowGraph cfg = new ControlFlowGraph(func.getBody());
		Map<Node, Knowledge> knowledge = calculateKnowledge(cfg);
		rewriteCode(cfg, knowledge);
	}

	private void rewriteCode(ControlFlowGraph cfg, Map<Node, Knowledge> knowledge) {
		for (Node node : cfg.getNodes()) {
			ImStmt stmt = node.getStmt();
			if (stmt == null) {
				continue;
			}
			Knowledge kn = knowledge.get(node);
			stmt.accept(new ImStmt.DefaultVisitor() {
				@Override
				public void visit(ImVarAccess va) {
					Value val = kn.varKnowledge.get(va.getVar());
					if (val == null) {
						return;
					}
					if (val.constantValue != null) {
						va.replaceWith(val.constantValue.copy());
					} else if (val.copyVar != null) {
						va.setVar(val.copyVar);
					}
				}
			});
			
		}
	}

	private Map<Node, Knowledge> calculateKnowledge(ControlFlowGraph cfg) {
		Map<Node, Knowledge> knowledge = new HashMap<>();

		// initialize with empty knowledge:<commit_msg>improved constant and copy propagation<commit_after>package de.peeeq.wurstscript.intermediateLang.optimizer;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.eclipse.jdt.annotation.NonNull;
import org.eclipse.jdt.annotation.Nullable;

import de.peeeq.wurstscript.intermediateLang.optimizer.ControlFlowGraph.Node;
import de.peeeq.wurstscript.jassIm.ImAlloc;
import de.peeeq.wurstscript.jassIm.ImConst;
import de.peeeq.wurstscript.jassIm.ImFunction;
import de.peeeq.wurstscript.jassIm.ImProg;
import de.peeeq.wurstscript.jassIm.ImSet;
import de.peeeq.wurstscript.jassIm.ImStmt;
import de.peeeq.wurstscript.jassIm.ImVar;
import de.peeeq.wurstscript.jassIm.ImVarAccess;
import de.peeeq.wurstscript.translation.imtranslation.ImTranslator;

public class ConstantAndCopyPropagation {

	private final ImProg prog;
	private final ImTranslator trans;

	public ConstantAndCopyPropagation(ImTranslator trans) {
		this.prog = trans.getImProg();
		this.trans = trans;
	}

	public void optimize() {
		for (ImFunction func : prog.getFunctions()) {
			optimizeFunc(func);
		}
	}

	static class Value {
		// one of the two is null
		final @Nullable ImVar copyVar;
		final @Nullable ImConst constantValue;

		public Value(ImVar copyVar) {
			this.copyVar = copyVar;
			this.constantValue = null;
		}

		public Value(ImConst constantValue) {
			this.copyVar = null;
			this.constantValue = constantValue;
		}
		
		@Override
		public boolean equals(@Nullable Object obj) {
			if (obj instanceof Value) {
				return equalValue((Value) obj);
			}
			return false;
		}

		public boolean equalValue(Value other) {
			if (copyVar == other.copyVar) {
				return true;
			} else if (constantValue != null && other.constantValue != null) {
				return constantValue.equalValue(other.constantValue);
			}
			return false;
		}
		
		@Override
		public String toString() {
			if (copyVar != null) {
				return "copy of " + copyVar;
			} else {
				return "constant " + constantValue;
			}
		}

	}

	static class Knowledge {

		Map<ImVar, Value> varKnowledge = new HashMap<>();
		Map<ImVar, Value> varKnowledgeOut = new HashMap<>();
		@Override
		public String toString() {
			return "[in =" + varKnowledge + ", out=" + varKnowledgeOut + "]";
		}
		
		

	}

	private void optimizeFunc(ImFunction func) {
		ControlFlowGraph cfg = new ControlFlowGraph(func.getBody());
		Map<Node, Knowledge> knowledge = calculateKnowledge(cfg);
		rewriteCode(cfg, knowledge);
	}

	private void rewriteCode(ControlFlowGraph cfg, Map<Node, Knowledge> knowledge) {
		for (Node node : cfg.getNodes()) {
			ImStmt stmt = node.getStmt();
			if (stmt == null) {
				continue;
			}
			Knowledge kn = knowledge.get(node);
			stmt.accept(new ImStmt.DefaultVisitor() {
				@Override
				public void visit(ImVarAccess va) {
					Value val = kn.varKnowledge.get(va.getVar());
					if (val == null) {
						return;
					}
					if (val.constantValue != null) {
						va.replaceWith(val.constantValue.copy());
					} else if (val.copyVar != null) {
						va.setVar(val.copyVar);
						// recursive call, because maybe it is possible to also replace the new var
						visit(va);
					}
				}
			});
			
		}
	}

	private Map<Node, Knowledge> calculateKnowledge(ControlFlowGraph cfg) {
		Map<Node, Knowledge> knowledge = new HashMap<>();

		// initialize with empty knowledge:
{'fraction_non_alphanumeric': 0.11586246749494365, 'fraction_numerical': 0.0, 'mean_word_length': 3.794321329639889, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4050368', 'n_tokens_mistral': 2636, 'n_tokens_neox': 2284, 'n_words': 509}