<commit_before>package org.ucombinator.jaam.visualizer.taint;

import javafx.scene.paint.Color;
import org.ucombinator.jaam.tools.taint3.Address;
import org.ucombinator.jaam.tools.taint3.Address.StaticField;
import org.ucombinator.jaam.tools.taint3.Address.InstanceField;
import org.ucombinator.jaam.tools.taint3.Address.ArrayRef;
import soot.SootClass;
import soot.SootField;
import soot.SootMethod;
import soot.Type;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;

public class TaintAddress extends TaintVertex {

    private static final Color defaultColor = Color.LIGHTGREEN;
    private static final Color primArrayRefColor = Color.SEAGREEN;
    private static final Color classArrayRefColor = Color.DARKGREEN;

    private Address address;
    private final boolean isField;
    private final String fieldId;
    private final SootClass sootClass;
    private final SootMethod sootMethod;

    public TaintAddress(Address address) {
        super(address.toString(), VertexType.TAINT_ADDRESS, true);
        this.color = defaultColor;
        this.address = address;

        this.isField = address instanceof StaticField || address instanceof InstanceField;
        if (address instanceof StaticField) {
            SootField f = ((StaticField) address).sootField();

            fieldId = address.sootClass().getName() + ":" + f.getName();
        }
        else if (address instanceof InstanceField) {
            SootField f = ((InstanceField) address).sootField();
            fieldId = address.sootClass().getName() + ":" + f.getName();
        }
        else {
           fieldId = null;
        }

        if (address instanceof ArrayRef) {
            soot.Type a = ((ArrayRef)address).typ();
            if (a instanceof soot.PrimType) {
                this.color = primArrayRefColor;
            }
            else {
                this.color = classArrayRefColor;
            }
        }

        sootClass = address.sootClass(); // Might be null
        sootMethod = address.sootMethod();
    }

    public TaintAddress copy() {
        TaintAddress newCopy =  new TaintAddress(this.address);
        newCopy.setColor(this.color);
        return newCopy;
    }

    public HashSet<String> getMethodNames() {
        if(address.sootMethod() != null) {
            HashSet<String> result = new HashSet<>();
            result.add(address.sootMethod().getName());
            return result;
        }
        else {
            return new HashSet<>();
        }
    }

    public Address getAddress() { return this.address; }

    public String getFieldId() { return fieldId; }

    public SootClass getSootClass() {
        return sootClass;
    }

    public SootMethod getSootMethod() {
        return sootMethod;
    }

    public boolean isArrayRef() {
        return address instanceof ArrayRef;
    }

    @Override
    public String getLongText() {
        return "Taint address:\n" + toString();
    }

    @Override
    public String toString() {
        return this.address.toString();
    }

    @Override
    public boolean hasField() { return this.isField; }

    @Override
    public void getFields(Collection<TaintAddress> store) {
        if(isField)
            store.add(this);
    }

    @Override
    public String getStmtString() {
        if(this.address.stmt() != null) {
            return this.address.stmt().toString();
        } else {
            return null;
        }
    }

    @Override
    public String getClassName() {
        if (sootClass == null) {
            return null;
        }
        else {
            return sootClass.getName();
        }
    }

    @Override
    public String getMethodName() {
        if (sootMethod != null) {
            return sootMethod.getName();
        }
        else {
            return null;
        }
    }

    @Override
    public List<TaintVertex> expand() {
        List<TaintVertex> expandedVertices = new ArrayList<>();
        expandedVertices.add(this);
        return expandedVertices;
    }
}<commit_msg>Improved TaintAddress and added types<commit_after>package org.ucombinator.jaam.visualizer.taint;

import javafx.scene.paint.Color;
import org.ucombinator.jaam.tools.taint3.Address;
import org.ucombinator.jaam.tools.taint3.Address.Return;
import org.ucombinator.jaam.tools.taint3.Address.Parameter;
import org.ucombinator.jaam.tools.taint3.Address.Throws;
import org.ucombinator.jaam.tools.taint3.Address.This;
import org.ucombinator.jaam.tools.taint3.Address.Lambda;
import org.ucombinator.jaam.tools.taint3.Address.StaticField;
import org.ucombinator.jaam.tools.taint3.Address.InstanceField;
import org.ucombinator.jaam.tools.taint3.Address.ArrayRef;
import org.ucombinator.jaam.visualizer.layout.LayoutAlgorithm;
import soot.SootClass;
import soot.SootField;
import soot.SootMethod;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;

public class TaintAddress extends TaintVertex {

    // Taken from Taint3, we summarixe only the one we care about
    public enum Type {
        Return, Parameter, Throws,
        Inner, // Any of: Stmt, Value, Local, New, NewArray, NewMultiArray,
        This, Lambda,
        StaticField, InstanceField,
        // The next two we split between primitive and class types
        ArrayRefPrim, ArrayRefClass
    }

    private Address address;
    // This is for the click on a field and see its graph functionality. It should probably
    // be improved by having a new fieldType thta inherits from address.. maybe later
    private final String fieldId;
    private final SootClass sootClass;
    private final SootMethod sootMethod;
    private final Type type;

    public TaintAddress(Address address) {
        super(address.toString(), VertexType.TAINT_ADDRESS, true);
        this.address = address;
        this.type = extractType();
        this.fieldId = extractFieldId();
        this.color = defaultColor();

        sootClass = address.sootClass(); // Might be null
        sootMethod = address.sootMethod();
    }


    public TaintAddress copy() {
        TaintAddress newCopy =  new TaintAddress(this.address);
        newCopy.setColor(this.color);
        return newCopy;
    }

    public HashSet<String> getMethodNames() {
        if(address.sootMethod() != null) {
            HashSet<String> result = new HashSet<>();
            result.add(address.sootMethod().getName());
            return result;
        }
        else {
            return new HashSet<>();
        }
    }

    public Address getAddress() { return this.address; }

    public String getFieldId() { return fieldId; }

    public SootClass getSootClass() {
        return sootClass;
    }

    public SootMethod getSootMethod() {
        return sootMethod;
    }

    public boolean isArrayRef() {
        return address instanceof ArrayRef;
    }

    @Override
    public String getLongText() {
        return "Taint address:\n" + toString();
    }

    @Override
    public String toString() {
        return this.address.toString();
    }

    private boolean isField() {
        return type == Type.InstanceField || type == Type.StaticField;
    }

    @Override
    public boolean hasField() { return isField(); }

    @Override
    public void getFields(Collection<TaintAddress> store) {
        if(isField())
            store.add(this);
    }

    @Override
    public String getStmtString() {
        if(this.address.stmt() != null) {
            return this.address.stmt().toString();
        } else {
            return null;
        }
    }

    @Override
    public String getClassName() {
        if (sootClass == null) {
            return null;
        }
        else {
            return sootClass.getName();
        }
    }

    @Override
    public String getMethodName() {
        if (sootMethod != null) {
            return sootMethod.getName();
        }
        else {
            return null;
        }
    }

    @Override
    public List<TaintVertex> expand() {
        List<TaintVertex> expandedVertices = new ArrayList<>();
        expandedVertices.add(this);
        return expandedVertices;
    }

    // No preferred layout
    @Override
    public LayoutAlgorithm.LAYOUT_ALGORITHM getPreferredLayout() {
        return null;
    }

    private Type extractType() {

        if (address instanceof Return) { return Type.Return; }
        if (address instanceof Parameter) { return Type.Return; }
        if (address instanceof Throws) { return Type.Return; }
        if (address instanceof This) { return Type.Return; }
        if (address instanceof Lambda) { return Type.Return; }
        if (address instanceof StaticField) { return Type.StaticField; }
        if (address instanceof InstanceField) { return Type.InstanceField; }
        if (address instanceof ArrayRef) {
            soot.Type a = ((ArrayRef)address).typ();
            if (a instanceof soot.PrimType) { return Type.ArrayRefPrim; }
            else { return Type.ArrayRefClass; }
        }
        return Type.Inner;
    }

    private String extractFieldId() {
        if (type != Type.StaticField || type != Type.InstanceField) return null;

        String name;
        if (type == Type.StaticField) {
            name = ((StaticField)address).sootField().getName();
        }
        else {
            name = ((InstanceField)address).sootField().getName();
        }
        return address.sootClass().getName() + ":" + name;
    }

    private Color defaultColor() {
        switch (type) {
            case Return: case Parameter: case Throws: return Color.LIGHTGREEN;
            case Inner: return Color.BLUE;
            case This: return Color.HOTPINK;
            case Lambda: return Color.ORCHID;
            case StaticField: return Color.BROWN;
            case InstanceField: return Color.SADDLEBROWN;
            case ArrayRefPrim: return Color.LIGHTGRAY;
            case ArrayRefClass: return Color.DARKGRAY;
        }
        // Shouldn't happen
        return Color.BLACK;
    }

}
{'fraction_non_alphanumeric': 0.09855936413313462, 'fraction_numerical': 0.0013909587680079483, 'mean_word_length': 2.319920844327177, 'pattern_counts': {'":': 3, '<': 19, '<?xml version=': 0, '>': 19, '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': '1566184', 'n_tokens_mistral': 2859, 'n_tokens_neox': 2768, 'n_words': 694}