<commit_before>package org.yinwang.pysonar.visitor;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.yinwang.pysonar.$;
import org.yinwang.pysonar.Analyzer;
import org.yinwang.pysonar.Binding;
import org.yinwang.pysonar.Builtins;
import org.yinwang.pysonar.State;
import org.yinwang.pysonar.ast.*;
import org.yinwang.pysonar.types.ClassType;
import org.yinwang.pysonar.types.DictType;
import org.yinwang.pysonar.types.FunType;
import org.yinwang.pysonar.types.InstanceType;
import org.yinwang.pysonar.types.ListType;
import org.yinwang.pysonar.types.ModuleType;
import org.yinwang.pysonar.types.TupleType;
import org.yinwang.pysonar.types.Type;
import org.yinwang.pysonar.types.UnionType;

import static org.yinwang.pysonar.Binding.Kind.ATTRIBUTE;
import static org.yinwang.pysonar.Binding.Kind.CLASS;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TypeInferencer implements Visitor1<Type, State> {

    @NotNull
    @Override
    public Type visit(Alias node, State s) {
        return Type.UNKNOWN;
    }

    @NotNull
    @Override
    public Type visit(Assert node, State s) {
        if (node.test != null) {
            visit(node.test, s);
        }
        if (node.msg != null) {
            visit(node.msg, s);
        }
        return Type.CONT;
    }

    @NotNull
    @Override
    public Type visit(Assign node, State s) {
        Type valueType = visit(node.value, s);
        bind(s, node.target, valueType);
        return Type.CONT;
    }

    @NotNull
    @Override
    public Type visit(Attribute node, State s) {
        Type targetType = visit(node.target, s);
        if (targetType instanceof UnionType) {
            Set<Type> types = ((UnionType) targetType).types;
            Type retType = Type.UNKNOWN;
            for (Type tt : types) {
                retType = UnionType.union(retType, getAttrType(node, tt));
            }
            return retType;
        } else {
            return getAttrType(node, targetType);
        }
    }

    @NotNull
    @Override
    public Type visit(Await node, State s) {
        if (node.value == null) {
            return Type.NONE;
        } else {
            return visit(node.value, s);
        }
    }

    @NotNull
    @Override
    public Type visit(BinOp node, State s) {
        Type ltype = visit(node.left, s);
        Type rtype = visit(node.right, s);
        Type result = applyOp(node.op, ltype, rtype, node.op.getMethod(), node, node.left);
        if (result != null) {
            return result;
        } else {
            Analyzer.self.putProblem(node, "Cannot apply binary operator " + node.op.getRep() +
                                           " to type " + ltype + " and " + rtype);
            return Type.UNKNOWN;
        }
    }

    @Nullable
    private Type applyOp(Op op, Type ltype, Type rtype, String method, Node node, Node left) {
        if (ltype instanceof InstanceType) {
            Type opType = ltype.table.lookupAttrType(method);
            if (opType instanceof FunType) {
                ((FunType) opType).setSelfType(ltype);
                return apply((FunType) opType, Collections.singletonList(rtype), null, null, null, node);
            } else {
                Analyzer.self.putProblem(left, "Operator method " + method + " is not a function");
                return null;
            }
        } else {
            return null;
        }
    }

    @NotNull
    @Override
    public Type visit(Block node, State s) {
        // first pass: mark global names
        for (Node n : node.seq) {
            if (n instanceof Global) {
                for (Name name : ((Global) n).names) {
                    s.addGlobalName(name.id);
                    Set<Binding> nb = s.lookup(name.id);
                    if (nb != null) {
                        Analyzer.self.putRef(name, nb);
                    }<commit_msg>use original method if operator is not overloaded<commit_after>package org.yinwang.pysonar.visitor;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.yinwang.pysonar.$;
import org.yinwang.pysonar.Analyzer;
import org.yinwang.pysonar.Binding;
import org.yinwang.pysonar.Builtins;
import org.yinwang.pysonar.State;
import org.yinwang.pysonar.ast.*;
import org.yinwang.pysonar.types.ClassType;
import org.yinwang.pysonar.types.DictType;
import org.yinwang.pysonar.types.FunType;
import org.yinwang.pysonar.types.InstanceType;
import org.yinwang.pysonar.types.ListType;
import org.yinwang.pysonar.types.ModuleType;
import org.yinwang.pysonar.types.TupleType;
import org.yinwang.pysonar.types.Type;
import org.yinwang.pysonar.types.UnionType;

import static org.yinwang.pysonar.Binding.Kind.ATTRIBUTE;
import static org.yinwang.pysonar.Binding.Kind.CLASS;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TypeInferencer implements Visitor1<Type, State> {

    @NotNull
    @Override
    public Type visit(Alias node, State s) {
        return Type.UNKNOWN;
    }

    @NotNull
    @Override
    public Type visit(Assert node, State s) {
        if (node.test != null) {
            visit(node.test, s);
        }
        if (node.msg != null) {
            visit(node.msg, s);
        }
        return Type.CONT;
    }

    @NotNull
    @Override
    public Type visit(Assign node, State s) {
        Type valueType = visit(node.value, s);
        bind(s, node.target, valueType);
        return Type.CONT;
    }

    @NotNull
    @Override
    public Type visit(Attribute node, State s) {
        Type targetType = visit(node.target, s);
        if (targetType instanceof UnionType) {
            Set<Type> types = ((UnionType) targetType).types;
            Type retType = Type.UNKNOWN;
            for (Type tt : types) {
                retType = UnionType.union(retType, getAttrType(node, tt));
            }
            return retType;
        } else {
            return getAttrType(node, targetType);
        }
    }

    @NotNull
    @Override
    public Type visit(Await node, State s) {
        if (node.value == null) {
            return Type.NONE;
        } else {
            return visit(node.value, s);
        }
    }

    @NotNull
    @Override
    public Type visit(BinOp node, State s) {
        Type ltype = visit(node.left, s);
        Type rtype = visit(node.right, s);
        if (operatorOverridden(ltype, node.op.getMethod())) {
            Type result = applyOp(node.op, ltype, rtype, node.op.getMethod(), node, node.left);
            if (result != null) {
                return result;
            }
        } else if (Op.isBoolean(node.op)) {
            return Type.BOOL;
        } else if (ltype == Type.UNKNOWN) {
            return rtype;
        } else if (rtype == Type.UNKNOWN) {
            return ltype;
        } else if (ltype.typeEquals(rtype)) {
            return ltype;
        }

        Analyzer.self.putProblem(node, "Cannot apply binary operator " + node.op.getRep() +
                                       " to type " + ltype + " and " + rtype);
        return Type.UNKNOWN;
    }

    private boolean operatorOverridden(Type type, String method) {
        if (type instanceof InstanceType) {
            Type opType = type.table.lookupAttrType(method);
            if (opType != null) {
                return true;
            }
        }
        return false;
    }

    @Nullable
    private Type applyOp(Op op, Type ltype, Type rtype, String method, Node node, Node left) {
        Type opType = ltype.table.lookupAttrType(method);
        if (opType instanceof FunType) {
            ((FunType) opType).setSelfType(ltype);
            return apply((FunType) opType, Collections.singletonList(rtype), null, null, null, node);
        } else {
            Analyzer.self.putProblem(left, "Operator method " + method + " is not a function");
            return null;
        }
    }

    @NotNull
    @Override
    public Type visit(Block node, State s) {
        // first pass: mark global names
        for (Node n : node.seq) {
            if (n instanceof Global) {
                for (Name name : ((Global) n).names) {
                    s.addGlobalName(name.id);
                    Set<Binding> nb = s.lookup(name.id);
                    if (nb != null) {
                        Analyzer.self.putRef(name, nb);
                    }
{'fraction_non_alphanumeric': 0.10882624768946396, 'fraction_numerical': 0.0002310536044362292, 'mean_word_length': 2.0665958200495926, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1790814', 'n_tokens_mistral': 2714, 'n_tokens_neox': 2602, 'n_words': 630}