path: root/config-model/src/main/java/com/yahoo/searchdefinition/MapEvaluationTypeContext.java

// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.searchdefinition;

import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
import com.yahoo.searchlib.rankingexpression.RankingExpression;
import com.yahoo.searchlib.rankingexpression.Reference;
import com.yahoo.searchlib.rankingexpression.parser.ParseException;
import com.yahoo.searchlib.rankingexpression.rule.Arguments;
import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
import com.yahoo.searchlib.rankingexpression.rule.FunctionReferenceContext;
import com.yahoo.searchlib.rankingexpression.rule.NameNode;
import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
import com.yahoo.tensor.TensorType;
import com.yahoo.tensor.evaluation.TypeContext;

import java.util.ArrayDeque;
import java.util.Collection;
import java.util.Collections;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.stream.Collectors;

/**
 * A context which only contains type information.
 * This returns empty tensor types (double) for unknown features which are not
 * query, attribute or constant features, as we do not have information about which such
 * features exist (but we know those that exist are doubles).
 *
 * This is not multithread safe.
 *
 * @author bratseth
 */
public class MapEvaluationTypeContext extends FunctionReferenceContext implements TypeContext<Reference> {

    private final Map<Reference, TensorType> featureTypes = new HashMap<>();

    private final Map<Reference, TensorType> resolvedTypes = new HashMap<>();

    /** For invocation loop detection */
    private final Deque<Reference> currentResolutionCallStack;

    private final SortedSet<Reference> queryFeaturesNotDeclared;
    private boolean tensorsAreUsed;

    private final MapEvaluationTypeContext parent;

    MapEvaluationTypeContext(Collection<ExpressionFunction> functions, Map<Reference, TensorType> featureTypes) {
        super(functions);
        this.featureTypes.putAll(featureTypes);
        this.currentResolutionCallStack =  new ArrayDeque<>();
        this.queryFeaturesNotDeclared = new TreeSet<>();
        tensorsAreUsed = false;
        parent = null;
    }

    private MapEvaluationTypeContext(Map<String, ExpressionFunction> functions,
                                     Map<String, String> bindings,
                                     Map<Reference, TensorType> featureTypes,
                                     Deque<Reference> currentResolutionCallStack,
                                     SortedSet<Reference> queryFeaturesNotDeclared,
                                     boolean tensorsAreUsed,
                                     MapEvaluationTypeContext parent) {
        super(functions, bindings);
        this.featureTypes.putAll(featureTypes);
        this.currentResolutionCallStack = currentResolutionCallStack;
        this.queryFeaturesNotDeclared = queryFeaturesNotDeclared;
        this.tensorsAreUsed = tensorsAreUsed;
        this.parent = parent;
    }

    public void setType(Reference reference, TensorType type) {
        featureTypes.put(reference, type);
        queryFeaturesNotDeclared.remove(reference);
    }

    @Override
    public TensorType getType(String reference) {
        throw new UnsupportedOperationException("Not able to parse general references from string form");
    }

    public void forgetResolvedTypes() {
        resolvedTypes.clear();
    }

    private TensorType resolvedType(Reference reference, int depth) {
//        System.out.println(indent + "In resolvedtype - resolving type for " + reference.toString());
        TensorType resolvedType = resolvedTypes.get(reference);
        if (resolvedType != null) {
//            System.out.println("Found previously resolved type for " + reference + " at depth " + depth + ": (" + resolvedType + ")");
            return resolvedType;
        }
        if (parent != null) return parent.resolvedType(reference, depth + 1);  // what about argument types? Careful with this!
//        System.out.println("Could NOT find type for " + reference + " - down to depth " + depth);
        return null;
    }

    private MapEvaluationTypeContext findOriginalParent() {
        if (parent != null)
            return parent.findOriginalParent();
        return this;
    }

    @Override
    public TensorType getType(Reference reference) {
        // computeIfAbsent without concurrent modification due to resolve adding more resolved entries:
        // TensorType resolvedType = resolvedTypes.get(reference);
        TensorType resolvedType = resolvedType(reference, 0);
        if (resolvedType != null) return resolvedType;

        resolvedType = resolveType(reference);
        if (resolvedType == null)
            return defaultTypeOf(reference); // Don't store fallback to default as we may know more later

//        System.out.println("Resolved type of " + reference + ": (" + resolvedType + ")");

        // Må inn her med et konsept av global eller lokal.
        // For globale - legg i lavest parent!
        MapEvaluationTypeContext originalParent = findOriginalParent();
        if (originalParent == null) {
            originalParent = this;
        }
        originalParent.resolvedTypes.put(reference, resolvedType);

        if (resolvedType.rank() > 0)
            tensorsAreUsed = true;
        return resolvedType;
    }

    private TensorType resolveType(Reference reference) {
        if (currentResolutionCallStack.contains(reference))
            throw new IllegalArgumentException("Invocation loop: " +
                                               currentResolutionCallStack.stream().map(Reference::toString).collect(Collectors.joining(" -> ")) +
                                               " -> " + reference);


        // Bound to a function argument, and not to a same-named identifier (which would lead to a loop)?
        Optional<String> binding = boundIdentifier(reference);
        if (binding.isPresent() && ! binding.get().equals(reference.toString())) {
            try {
                // This is not pretty, but changing to bind expressions rather
                // than their string values requires deeper changes
                return new RankingExpression(binding.get()).type(this);
            } catch (ParseException e) {
                throw new IllegalArgumentException(e);
            }
        }

        try {
            currentResolutionCallStack.addLast(reference);

            // A reference to an attribute, query or constant feature?
            if (FeatureNames.isSimpleFeature(reference)) {
                // The argument may be a local identifier bound to the actual value
                String argument = reference.simpleArgument().get();
                String argumentBinding = getBinding(argument);
                reference = Reference.simple(reference.name(), argumentBinding != null ? argumentBinding : argument);
                return featureTypes.get(reference);
            }

            // A reference to a function?
            Optional<ExpressionFunction> function = functionInvocation(reference);
            if (function.isPresent()) {
                return function.get().getBody().type(this.withBindings(bind(function.get().arguments(), reference.arguments())));
            }

            // A reference to a feature which returns a tensor?
            Optional<TensorType> featureTensorType = tensorFeatureType(reference);
            if (featureTensorType.isPresent()) {
                return featureTensorType.get();
            }

            // A directly injected identifier? (Useful for stateless model evaluation)
            if (reference.isIdentifier() && featureTypes.containsKey(reference)) {
                return featureTypes.get(reference);
            }

            // We do not know what this is - since we do not have complete knowledge about the match features
            // in Java we must assume this is a match feature and return the double type - which is the type of
            // all match features
            return TensorType.empty;
        }
        finally {
            currentResolutionCallStack.removeLast();
        }
    }

    /**
     * Returns the default type for this simple feature, or null if it does not have a default
     */
    public TensorType defaultTypeOf(Reference reference) {
        if ( ! FeatureNames.isSimpleFeature(reference))
            throw new IllegalArgumentException("This can only be called for simple references, not " + reference);
        if (reference.name().equals("query")) { // we do not require all query features to be declared, only non-doubles
            queryFeaturesNotDeclared.add(reference);
            return TensorType.empty;
        }
        return null;
    }

    /**
     * Returns the binding if this reference is a simple identifier which is bound in this context.
     * Returns empty otherwise.
     */
    private Optional<String> boundIdentifier(Reference reference) {
        if ( ! reference.arguments().isEmpty()) return Optional.empty();
        if ( reference.output() != null) return Optional.empty();
        return Optional.ofNullable(getBinding(reference.name()));
    }

    private Optional<ExpressionFunction> functionInvocation(Reference reference) {
        if (reference.output() != null) return Optional.empty();
        ExpressionFunction function = functions().get(reference.name());
        if (function == null) return Optional.empty();
        if (function.arguments().size() != reference.arguments().size()) return Optional.empty();
        return Optional.of(function);
    }

    /**
     * There are two features which returns the (non-empty) tensor type: tensorFromLabels and tensorFromWeightedSet.
     * This returns the type of those features if this is a reference to either of them, or empty otherwise.
     */
    private Optional<TensorType> tensorFeatureType(Reference reference) {
        if ( ! reference.name().equals("tensorFromLabels") && ! reference.name().equals("tensorFromWeightedSet"))
            return Optional.empty();

        if (reference.arguments().size() != 1 && reference.arguments().size() != 2)
            throw new IllegalArgumentException(reference.name() + " must have one or two arguments");

        ExpressionNode arg0 = reference.arguments().expressions().get(0);
        if ( ! ( arg0 instanceof ReferenceNode) || ! FeatureNames.isSimpleFeature(((ReferenceNode)arg0).reference()))
            throw new IllegalArgumentException("The first argument of " + reference.name() +
                                               " must be a simple feature, not " + arg0);

        String dimension;
        if (reference.arguments().size() > 1) {
            ExpressionNode arg1 = reference.arguments().expressions().get(1);
            if ( ( ! (arg1 instanceof ReferenceNode) || ! (((ReferenceNode)arg1).reference().isIdentifier()))
                 &&
                 ( ! (arg1 instanceof NameNode)))
                throw new IllegalArgumentException("The second argument of " + reference.name() +
                                                   " must be a dimension name, not " + arg1);
            dimension = reference.arguments().expressions().get(1).toString();
        }
        else { // default
            dimension = ((ReferenceNode)arg0).reference().arguments().expressions().get(0).toString();
        }

        // TODO: Determine the type of the weighted set/vector and use that as value type
        return Optional.of(new TensorType.Builder().mapped(dimension).build());
    }

    /** Binds the given list of formal arguments to their actual values */
    private Map<String, String> bind(List<String> formalArguments,
                                     Arguments invocationArguments) {
        Map<String, String> bindings = new HashMap<>(formalArguments.size());
        for (int i = 0; i < formalArguments.size(); i++) {
            String identifier = invocationArguments.expressions().get(i).toString();
            String identifierBinding = super.getBinding(identifier);
            bindings.put(formalArguments.get(i), identifierBinding != null ? identifierBinding : identifier);
        }
        return bindings;
    }

    public Map<Reference, TensorType> featureTypes() {
        return Collections.unmodifiableMap(featureTypes);
    }

    /**
     * Returns an unmodifiable view of the query features which was requested but for which we have no type info
     * (such that they default to TensorType.empty), shared between all instances of this
     * involved in resolving a particular rank profile.
     */
    public SortedSet<Reference> queryFeaturesNotDeclared() {
        return Collections.unmodifiableSortedSet(queryFeaturesNotDeclared);
    }

    /** Returns true if any feature across all instances involved in resolving this rank profile resolves to a tensor */
    public boolean tensorsAreUsed() { return tensorsAreUsed; }

    @Override
    public MapEvaluationTypeContext withBindings(Map<String, String> bindings) {
        return new MapEvaluationTypeContext(functions(),
                                            bindings,
                                            featureTypes,
                                            currentResolutionCallStack,
                                            queryFeaturesNotDeclared,
                                            tensorsAreUsed, this);
    }

}