package org.broadinstitute.hellbender.tools.walkers.varianteval;

import htsjdk.samtools.SAMSequenceRecord;
import htsjdk.samtools.reference.IndexedFastaSequenceFile;
import htsjdk.samtools.reference.ReferenceSequenceFile;
import htsjdk.tribble.Feature;
import htsjdk.variant.variantcontext.Allele;
import htsjdk.variant.variantcontext.VariantContext;
import htsjdk.variant.variantcontext.VariantContextBuilder;
import htsjdk.variant.variantcontext.VariantContextUtils;
import htsjdk.variant.vcf.VCFHeader;
import org.apache.logging.log4j.Logger;
import org.broadinstitute.barclay.argparser.*;
import org.broadinstitute.barclay.help.DocumentedFeature;
import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;
import org.broadinstitute.hellbender.cmdline.argumentcollections.DbsnpArgumentCollection;
import org.broadinstitute.hellbender.cmdline.argumentcollections.MultiVariantInputArgumentCollection;
import org.broadinstitute.hellbender.engine.*;
import org.broadinstitute.hellbender.exceptions.GATKException;
import org.broadinstitute.hellbender.exceptions.UserException;
import org.broadinstitute.hellbender.tools.walkers.varianteval.evaluators.*;
import org.broadinstitute.hellbender.tools.walkers.varianteval.stratifications.AlleleFrequency.*;
import org.broadinstitute.hellbender.tools.walkers.varianteval.stratifications.IntervalStratification;
import org.broadinstitute.hellbender.tools.walkers.varianteval.stratifications.VariantStratifier;
import org.broadinstitute.hellbender.tools.walkers.varianteval.stratifications.manager.StratificationManager;
import org.broadinstitute.hellbender.tools.walkers.varianteval.util.EvaluationContext;
import org.broadinstitute.hellbender.tools.walkers.varianteval.util.SortableJexlVCMatchExp;
import org.broadinstitute.hellbender.tools.walkers.varianteval.util.VariantEvalUtils;
import org.broadinstitute.hellbender.utils.SimpleInterval;
import org.broadinstitute.hellbender.utils.Utils;
import org.broadinstitute.hellbender.utils.io.IOUtils;
import org.broadinstitute.hellbender.utils.samples.PedigreeValidationType;
import org.broadinstitute.hellbender.utils.samples.SampleDB;
import org.broadinstitute.hellbender.utils.samples.SampleDBBuilder;
import org.broadinstitute.hellbender.utils.variant.GATKVariantContextUtils;
import picard.cmdline.programgroups.VariantEvaluationProgramGroup;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.PrintStream;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 *
 *
 *

 * Given a variant callset, it is common to calculate various quality control metrics. These metrics include the number of
 * raw or filtered SNP counts; ratio of transition mutations to transversions; concordance of a particular sample's calls
 * to a genotyping chip; number of s per sample; etc. Furthermore, it is often useful to stratify these metrics
 * by various criteria like functional class (missense, nonsense, silent), whether the site is CpG site, the amino acid
 * degeneracy of the site, etc. VariantEval facilitates these calculations in two ways: by providing several built-in
 * evaluation and stratification modules, and by providing a framework that permits the easy development of new evaluation
 * and stratification modules.
 *

 *
 *
Input

 *

 * One or more variant sets to evaluate plus any number of comparison sets.
 *

 *
 *
Output

 *

 * Evaluation tables detailing the results of the eval modules which were applied.
 * For example:
 *

 *

 * output.eval.grp:
 * ##:GATKReport.v0.1 CountVariants : Counts different classes of variants in the sample
 * CountVariants  CompFeatureInput  CpG      EvalFeatureInput  JexlExpression  Novelty  nProcessedLoci  nCalledLoci  nRefLoci  nVariantLoci  variantRate ...
 * CountVariants  dbsnp             CpG      eval              none            all      65900028        135770       0         135770        0.00206024  ...
 * CountVariants  dbsnp             CpG      eval              none            known    65900028        47068        0         47068         0.00071423  ...
 * CountVariants  dbsnp             CpG      eval              none            novel    65900028        88702        0         88702         0.00134601  ...
 * CountVariants  dbsnp             all      eval              none            all      65900028        330818       0         330818        0.00502000  ...
 * CountVariants  dbsnp             all      eval              none            known    65900028        120685       0         120685        0.00183133  ...
 * CountVariants  dbsnp             all      eval              none            novel    65900028        210133       0         210133        0.00318866  ...
 * CountVariants  dbsnp             non_CpG  eval              none            all      65900028        195048       0         195048        0.00295976  ...
 * CountVariants  dbsnp             non_CpG  eval              none            known    65900028        73617        0         73617         0.00111710  ...
 * CountVariants  dbsnp             non_CpG  eval              none            novel    65900028        121431       0         121431        0.00184265  ...
 * ...
 *

 *

 *
 *
Usage examples

 *

 * gatk VariantEval \
 *   -R reference.fasta \
 *   -O output.eval.grp \
 *   --eval set1:set1.vcf \
 *   --eval set2:set2.vcf \
 *   [--comp comp.vcf]
 *

 *
 * Count Mendelian violations for each family in a callset with multiple families (and provided pedigree)
 *

 * gatk VariantEval \
 *   -R reference.fasta \
 *   -O output.MVs.byFamily.table \
 *   --eval multiFamilyCallset.vcf \
 *   -no-ev -noST \
 *   -ST Family \
 *   -EV MendelianViolationEvaluator
 *

 *
 *
Caveat

 *
 *
Some stratifications and evaluators are incompatible with each other due to their respective memory requirements,
 * such as AlleleCount and VariantSummary, or Sample and VariantSummary. If you specify such a combination, the program
 * will output an error message and ask you to disable one of these options. We do not currently provide an exhaustive
 * list of incompatible combinations, so we recommend trying out combinations that you are interested in on a dummy
 * command line, to rapidly ascertain whether it will work or not.

 *
 */
@CommandLineProgramProperties(
        summary = "Given a variant callset, it is common to calculate various quality control metrics. These metrics include the number of " +
                "raw or filtered SNP counts; ratio of transition mutations to transversions; concordance of a particular sample's calls " +
                "to a genotyping chip; number of singletons per sample; etc. Furthermore, it is often useful to stratify these metrics " +
                "by various criteria like functional class (missense, nonsense, silent), whether the site is CpG site, the amino acid " +
                "degeneracy of the site, etc. VariantEval facilitates these calculations in two ways: by providing several built-in " +
                "evaluation and stratification modules, and by providing a framework that permits the easy development of new evaluation " +
                "and stratification modules.",
        oneLineSummary = "General-purpose tool for variant evaluation (% in dbSNP, genotype concordance, Ti/Tv ratios, and a lot more)",
        programGroup = VariantEvaluationProgramGroup.class
)
@DocumentedFeature
@BetaFeature
public class VariantEval extends MultiVariantWalker {
    public static final String IS_SINGLETON_KEY = "ISSINGLETON";

    @Argument(fullName = StandardArgumentDefinitions.OUTPUT_LONG_NAME,
            shortName = StandardArgumentDefinitions.OUTPUT_SHORT_NAME,
            doc="File to which variants should be written")
    protected File outFile;

    /**
     * The variant file(s) to evaluate.
     */
    @Argument(fullName="eval", shortName = "eval", doc="Input evaluation file(s)", optional=false)
    public List<FeatureInput<VariantContext>> evals;

    /**
     * The variant file(s) to compare against.
     */
    @Argument(fullName = StandardArgumentDefinitions.COMPARISON_LONG_NAME, shortName = StandardArgumentDefinitions.COMPARISON_SHORT_NAME, doc="Input comparison file(s)", optional=true)
    public List<FeatureInput<VariantContext>> compsProvided = new ArrayList<>();
    private List<FeatureInput<VariantContext>> comps = new ArrayList<>();

    /**
     * dbSNP comparison VCF.  By default, the dbSNP file is used to specify the set of "known" variants.
     * Other sets can be specified with the -known-name (--known_names) argument.
     */
    @ArgumentCollection
    protected DbsnpArgumentCollection dbsnp = new DbsnpArgumentCollection();

    /**
     * Some analyses want to count overlap not with dbSNP (which is in general very open) but
     * actually want to itemize their overlap specifically with a set of gold standard sites
     * such as HapMap, OMNI, or the gold standard indels.  This argument provides a mechanism
     * for communicating which file to use
     */
    @Argument(fullName="gold-standard", shortName = "gold", doc="Evaluations that count calls at sites of true variation (e.g., indel calls) will use this argument as their gold standard for comparison", optional=true)
    public FeatureInput<VariantContext> goldStandard = null;

    /**
     * Note that the --list argument requires a fully resolved and correct command-line to work.
     */
    @Argument(fullName="list", shortName="ls", doc="List the available eval modules and exit", optional=true)
    protected Boolean LIST = false;

    // Partitioning the data arguments
    @Argument(shortName="select", doc="One or more stratifications to use when evaluating the data", optional=true)
    protected ArrayList<String> SELECT_EXPS = new ArrayList<String>();

    @Argument(shortName="select-name", doc="Names to use for the list of stratifications (must be a 1-to-1 mapping)", optional=true)
    protected ArrayList<String> SELECT_NAMES = new ArrayList<String>();

    @Argument(fullName="sample", shortName="sn", doc="Derive eval and comp contexts using only these sample genotypes, when genotypes are available in the original context", optional=true)
    protected Set<String> SAMPLE_EXPRESSIONS = new TreeSet<>();

    @Argument(fullName = StandardArgumentDefinitions.PEDIGREE_FILE_LONG_NAME, shortName = StandardArgumentDefinitions.PEDIGREE_FILE_SHORT_NAME, doc="Pedigree file for determining the population \"founders\"", optional=true)
    private File pedigreeFile;

    /**
     * List of feature tracks to be used for specifying "known" variants other than dbSNP.
     */
    @Argument(shortName="known-name", doc="Name of feature bindings containing variant sites that should be treated as known when splitting eval features into known and novel subsets", optional=true)
    protected Set<String> KNOWN_NAMES = new HashSet<String>();
    List<FeatureInput<VariantContext>> knowns = new ArrayList<>();

    // Stratification arguments
    @Argument(fullName="stratification-module", shortName="ST", doc="One or more specific stratification modules to apply to the eval track(s) (in addition to the standard stratifications, unless -noS is specified)", optional=true)
    protected List<String> STRATIFICATIONS_TO_USE = new ArrayList<>();

    @Argument(fullName="do-not-use-all-standard-stratifications", shortName="no-st", doc="Do not use the standard stratification modules by default (instead, only those that are specified with the -S option)", optional=true)
    protected Boolean NO_STANDARD_STRATIFICATIONS = false;

    /**
     * See the -list argument to view available modules.
     */
    @Argument(fullName="eval-module", shortName="EV", doc="One or more specific eval modules to apply to the eval track(s) (in addition to the standard modules, unless -no-ev is specified)", optional=true)
    protected List<String> MODULES_TO_USE = new ArrayList<>();

    @Argument(fullName="do-not-use-all-standard-modules", shortName="no-ev", doc="Do not use the standard modules by default (instead, only those that are specified with the -EV option)", optional=true)
    protected Boolean NO_STANDARD_MODULES = false;

    @Argument(fullName="min-phase-quality", shortName="mpq", doc="Minimum phasing quality", optional=true)
    protected double MIN_PHASE_QUALITY = 10.0;

    @Argument(shortName="mvq", fullName="mendelian-violation-qual-threshold", doc="Minimum genotype QUAL score for each trio member required to accept a site as a violation. Default is 50.", optional=true)
    protected double MENDELIAN_VIOLATION_QUAL_THRESHOLD = 50;

    @Argument(shortName="ploidy", fullName="sample-ploidy", doc="Per-sample ploidy (number of chromosomes per sample)", optional=true)
    protected int ploidy = GATKVariantContextUtils.DEFAULT_PLOIDY;

    @Argument(fullName="ancestral-alignments", shortName="aa", doc="Fasta file with ancestral alleles", optional=true)
    private File ancestralAlignmentsFile = null;

    @Argument(fullName="require-strict-allele-match", shortName="strict", doc="If provided only comp and eval tracks with exactly matching reference and alternate alleles will be counted as overlapping", optional=true)
    private boolean requireStrictAlleleMatch = false;

    @Argument(fullName="keep-ac0", shortName="keep-ac0", doc="If provided, modules that track polymorphic sites will not require that a site have AC > 0 when the input eval has genotypes", optional=true)
    protected boolean keepSitesWithAC0 = false;

    @Hidden
    @Argument(fullName="num-samples", doc="If provided, modules that track polymorphic sites will not require that a site have AC > 0 when the input eval has genotypes", optional=true)
    private int numSamplesFromArgument = 0;

    /**
     * If true, VariantEval will treat -eval 1 -eval 2 as separate tracks from the same underlying
     * variant set, and evaluate the union of the results.  Useful when you want to do -eval chr1.vcf -eval chr2.vcf etc.
     */
    @Argument(fullName="merge-evals", shortName="merge-evals", doc="If provided, all -eval tracks will be merged into a single eval track", optional=true)
    public boolean mergeEvals = false;

    /**
     * File containing tribble-readable features for the IntervalStratificiation
     */
    @Argument(fullName="strat-intervals", shortName="strat-intervals", doc="File containing tribble-readable features for the IntervalStratificiation", optional=true)
    public FeatureInput<Feature> intervalsFile = null;

    /**
     * File containing tribble-readable features containing known CNVs.  For use with VariantSummary table.
     */
    @Argument(fullName="known-cnvs", shortName="known-cnvs", doc="File containing tribble-readable features describing a known list of copy number variants", optional=true)
    public FeatureInput<Feature> knownCNVsFile = null;

    protected StratifyingScale AFScale = StratifyingScale.LINEAR;
    protected boolean useCompAFStratifier = false;

    @Override
    protected MultiVariantInputArgumentCollection getMultiVariantInputArgumentCollection() {
        return new MultiVariantInputArgumentCollection() {
            private static final long serialVersionUID = 1L;

            @Override
            public List<String> getDrivingVariantPaths() {
                //driving variants will be determined by initializeDrivingVariants()
                return Collections.emptyList();
            }
        };
    }

    @Override
    protected void initializeDrivingVariants() {
        getDrivingVariantsFeatureInputs().addAll(evals);
        if (dbsnp.dbsnp != null) {
            getDrivingVariantsFeatureInputs().add(dbsnp.dbsnp);
        }

        getDrivingVariantsFeatureInputs().addAll(compsProvided);

        super.initializeDrivingVariants();
    }

    // Variables
    private Set<SortableJexlVCMatchExp> jexlExpressions = new TreeSet<SortableJexlVCMatchExp>();

    private boolean isSubsettingSamples;
    private Set<String> sampleNamesForEvaluation = new LinkedHashSet<String>();
    private Set<String> familyNamesForEvaluation = new LinkedHashSet<String>();
    private Set<String> sampleNamesForStratification = new LinkedHashSet<String>();
    private Set<String> familyNamesForStratification = new LinkedHashSet<String>();

    // important stratifications
    private boolean byFilterIsEnabled = false;
    private boolean perSampleIsEnabled = false;
    private boolean perFamilyIsEnabled = false;

    // Public constants
    final private static String ALL_SAMPLE_NAME = "all";
    final private static String ALL_FAMILY_NAME = "all";

    // Utility class
    private final VariantEvalUtils variantEvalUtils = new VariantEvalUtils(this);

    // Ancestral alignments
    private ReferenceSequenceFile ancestralAlignments = null;

    // The set of all possible evaluation contexts
    StratificationManager<VariantStratifier, EvaluationContext> stratManager;

    private SampleDB sampleDB = null;

    // maintain the mapping of FeatureInput to name used in output file
    Map<FeatureInput<VariantContext>, String> inputToNameMap = new HashMap<>();

    /**
     * Initialize the stratifications, evaluations, evaluation contexts, and reporting object
     */
    @Override
    public void onTraversalStart() {
        Utils.nonNull(outFile);

        // Just list the modules, and exit quickly.
        if (LIST) { variantEvalUtils.listModulesAndExit(); }

        sampleDB = initializeSampleDB();

        comps.addAll(compsProvided);
        compsProvided.forEach(x -> inputToNameMap.put(x, x.hasUserSuppliedName() ? x.getName() : StandardArgumentDefinitions.COMPARISON_SHORT_NAME));
        if ( dbsnp.dbsnp != null ) {
            comps.add(dbsnp.dbsnp);
            inputToNameMap.put(dbsnp.dbsnp, "dbsnp");
            knowns.add(dbsnp.dbsnp);
        }

        evals.forEach(x -> inputToNameMap.put(x, x.hasUserSuppliedName() ? x.getName() : "eval"));

        // Set up set of additional knowns
        for ( FeatureInput<VariantContext> compInput : comps ) {
            if ( KNOWN_NAMES.contains(getNameForInput(compInput)))
                knowns.add(compInput);
        }

        // Now that we have all the inputs categorized, determine the sample list from the eval inputs.
        Map<String, VCFHeader> vcfInputs = new HashMap<>();
        evals.forEach(x -> vcfInputs.put(x.getName(), (VCFHeader)getHeaderForFeatures(x)));

        Set<String> vcfSamples = new HashSet<>();
        vcfInputs.forEach((k,v) -> vcfSamples.addAll(v.getSampleNamesInOrder()));

        // Load the sample list, using an intermediate tree set to sort the samples
        final Set<String> allSampleNames = new HashSet<>(vcfSamples);
        sampleNamesForEvaluation.addAll(new TreeSet<>(SAMPLE_EXPRESSIONS.isEmpty() ? vcfSamples : Utils.filterCollectionByExpressions(vcfSamples, SAMPLE_EXPRESSIONS, false)));

        isSubsettingSamples = ! sampleNamesForEvaluation.containsAll(allSampleNames);
        familyNamesForEvaluation.addAll(sampleDB.getFamilyIDs());

        //If stratifying by sample name, assign a stratification for each sample we're evaluating (based on commandline args)...
        if (STRATIFICATIONS_TO_USE.contains("Sample") ) {
            sampleNamesForStratification.addAll(sampleNamesForEvaluation);
        }
        //...and also a stratification for the sum over all samples
        sampleNamesForStratification.add(ALL_SAMPLE_NAME);

        //If stratifying by sample name, assign a stratification for each family...
        if ( STRATIFICATIONS_TO_USE.contains("Family") ) {
            familyNamesForStratification.addAll(familyNamesForEvaluation);
        }
        //...and also a stratification for the sum over all families
        familyNamesForStratification.add(ALL_FAMILY_NAME);

        // Initialize select expressions
        for (VariantContextUtils.JexlVCMatchExp jexl : VariantContextUtils.initializeMatchExps(SELECT_NAMES, SELECT_EXPS)) {
            SortableJexlVCMatchExp sjexl = new SortableJexlVCMatchExp(jexl.name, jexl.exp);
            jexlExpressions.add(sjexl);
        }

        // Initialize the set of stratifications and evaluations to use
        // The list of stratifiers and evaluators to use
        final List<VariantStratifier> stratificationObjects = variantEvalUtils.initializeStratificationObjects(NO_STANDARD_STRATIFICATIONS, STRATIFICATIONS_TO_USE);
        final Set<Class<? extends VariantEvaluator>> evaluationClasses = variantEvalUtils.initializeEvaluationObjects(NO_STANDARD_MODULES, MODULES_TO_USE);

        checkForIncompatibleEvaluatorsAndStratifiers(stratificationObjects, evaluationClasses);

        for ( VariantStratifier vs : stratificationObjects ) {
            if ( vs.getName().equals("Filter") )
                byFilterIsEnabled = true;
            else if ( vs.getName().equals("Sample") )
                perSampleIsEnabled = true;
            else if ( vs.getName().equals("Family"))
                perFamilyIsEnabled = true;
        }

        if (perSampleIsEnabled && perFamilyIsEnabled)
            throw new CommandLineException.BadArgumentValue("ST", "Variants cannot be stratified by sample and family at the same time");

        if (perFamilyIsEnabled && sampleDB.getTrios().isEmpty())
            throw new CommandLineException.BadArgumentValue("ST", "Cannot stratify by family without *.ped file");

        if ( intervalsFile != null ) {
            boolean fail = true;
            for ( final VariantStratifier vs : stratificationObjects ) {
                if ( vs.getClass().equals(IntervalStratification.class) )
                    fail = false;
            }
            if ( fail )
                throw new CommandLineException.BadArgumentValue("ST", "stratIntervals argument provided but -ST IntervalStratification not provided");
        }

        // Initialize the evaluation contexts
        createStratificationStates(stratificationObjects, evaluationClasses);

        // Load ancestral alignments
        if (ancestralAlignmentsFile != null) {
            try {
                ancestralAlignments = new IndexedFastaSequenceFile(ancestralAlignmentsFile);
            } catch (FileNotFoundException e) {
                throw new GATKException(String.format("The ancestral alignments file, '%s', could not be found", ancestralAlignmentsFile.getAbsolutePath()));
            }
        }

        assertThatTerritoryIsSpecifiedIfNecessary();
    }

    private void assertThatTerritoryIsSpecifiedIfNecessary() {
        final Set<String> evaluatorsWhichRequireTerritory = stratManager.values()
                .stream()
                .flatMap(ctx -> ctx.getVariantEvaluators().stream())
                .filter(Objects::nonNull)
                .filter(VariantEvaluator::requiresTerritoryToBeSpecified)
                .map(VariantEvaluator::getSimpleName)
                .collect(Collectors.toSet());
        if(!evaluatorsWhichRequireTerritory.isEmpty() && getTraversalIntervals() == null){
            throw new UserException("You specified evaluators which require a covered territory to be specified.  " +
                    "\nPlease specify intervals or a reference file or disable all of the following evaluators:" +
                    evaluatorsWhichRequireTerritory.stream()
                            .collect(Collectors.joining(", ")));
        }
    }

    private void checkForIncompatibleEvaluatorsAndStratifiers( final List<VariantStratifier> stratificationObjects,
                                                             Set<Class<? extends VariantEvaluator>> evaluationClasses) {
        for ( final VariantStratifier vs : stratificationObjects ) {
            for ( Class<? extends VariantEvaluator> ec : evaluationClasses )
                if ( vs.getIncompatibleEvaluators().contains(ec) )
                    throw new CommandLineException.BadArgumentValue("ST and ET",
                            "The selected stratification " + vs.getName() +
                                    " and evaluator " + ec.getSimpleName() +
                                    " are incompatible due to combinatorial memory requirements." +
                                    " Please disable one");
        }
    }

    final void createStratificationStates(final List<VariantStratifier> stratificationObjects, final Set<Class<? extends VariantEvaluator>> evaluationObjects) {
        final List<VariantStratifier> strats = new ArrayList<VariantStratifier>(stratificationObjects);
        stratManager = new StratificationManager<>(strats);

        logger.info("Creating " + stratManager.size() + " combinatorial stratification states");
        for ( int i = 0; i < stratManager.size(); i++ ) {
            EvaluationContext ec = createEvaluationContext(evaluationObjects);
            stratManager.set(i, ec);
        }
    }

    /**
     * Create the EvaluationContext (new instance) for the provided set of VariantEvaluators.
     *
     * @param evaluationObjects The list of VariantEvaluator classes
     * @return The EvaluationContext for this set of VariantEvaluator classes
     */
    protected EvaluationContext createEvaluationContext(final Set<Class<? extends VariantEvaluator>> evaluationObjects) {
        return new EvaluationContext(this, evaluationObjects);
    }

    private class PositionAggregator {
        private SimpleInterval i = null;

        private ReadsContext readsContext;
        private ReferenceContext referenceContext;
        private FeatureContext featureContext;

        private void addVariant(VariantContext vc, ReadsContext readsContext, ReferenceContext referenceContext, FeatureContext featureContext) {
            if (i == null || !vc.getContig().equals(i.getContig()) || vc.getStart() != i.getStart()) {
                callDoApply();

                i = new SimpleInterval(vc.getContig(), vc.getStart(), vc.getEnd());
                this.readsContext = readsContext;
                this.referenceContext = referenceContext;
                this.featureContext = featureContext;
            }
            else if (vc.getEnd() > i.getEnd()) {
                //expand region
                i = new SimpleInterval(i.getContig(), i.getStart(), vc.getEnd());

                this.readsContext = new ReadsContext(this.readsContext, i);
                this.referenceContext = new ReferenceContext(this.referenceContext, i);
                this.featureContext = new FeatureContext(this.featureContext, i);
            }
        }

        public void callDoApply(){
            if (i != null) {
                doApply(this.readsContext, this.referenceContext, this.featureContext);
                i = null;
            }
        }

        public void onComplete() {
            callDoApply();
        }
    }

    final PositionAggregator aggr = new PositionAggregator();

    @Override
    public void apply(VariantContext variant, ReadsContext readsContext, ReferenceContext referenceContext, FeatureContext featureContext) {
        aggr.addVariant(variant, readsContext, referenceContext, featureContext);
    }

    public String getNameForInput(FeatureInput<VariantContext> input) {
        return inputToNameMap.get(input);
    }

    /**
     * This will get called once per site where a variant is present in any input
     */
    public void doApply(ReadsContext readsContext, ReferenceContext referenceContext, FeatureContext featureContext) {
        HashMap<FeatureInput<VariantContext>, HashMap<String, Collection<VariantContext>>> evalVCs = variantEvalUtils.bindVariantContexts(referenceContext, featureContext, evals, byFilterIsEnabled, true, perSampleIsEnabled, perFamilyIsEnabled, mergeEvals);
        HashMap<FeatureInput<VariantContext>, HashMap<String, Collection<VariantContext>>> compVCs = variantEvalUtils.bindVariantContexts(referenceContext, featureContext, comps, byFilterIsEnabled, false, false, false, false);

        // for each eval track
        for ( final FeatureInput<VariantContext> evalInput : evals ) {
            final Map<String, Collection<VariantContext>> emptyEvalMap = Collections.emptyMap();
            final Map<String, Collection<VariantContext>> evalSet = evalVCs.containsKey(evalInput) ? evalVCs.get(evalInput) : emptyEvalMap;

            Set<String> statificationLevels;

            // for each sample stratifier
            if (perFamilyIsEnabled)
                statificationLevels = familyNamesForStratification;
            else
                statificationLevels = sampleNamesForStratification;
            for ( final String stratLevelName : statificationLevels ) {
                Collection<VariantContext> evalSetBySample = evalSet.get(stratLevelName);

                if ( evalSetBySample == null ) {
                    evalSetBySample = new HashSet<>(1);
                    evalSetBySample.add(null);
                }

                // for each eval in the track
                for ( VariantContext eval : evalSetBySample ) {
                    String aastr = (ancestralAlignments == null) ? null : new String(ancestralAlignments.getSubsequenceAt(eval.getContig(), eval.getStart(), eval.getEnd()).getBases());

                    // deal with ancestral alleles if requested
                    if ( eval != null && aastr != null ) {
                        eval = new VariantContextBuilder(eval).attribute("ANCESTRALALLELE", aastr).make();
                    }

                    String evalName = getNameForInput(evalInput);

                    // for each comp track
                    for ( final FeatureInput<VariantContext> compInput : comps ) {
                        processComp(referenceContext, readsContext, featureContext, eval, evalName, compInput, stratLevelName, compVCs, evalSetBySample);
                    }

                    if (comps.isEmpty()) {
                        processComp(referenceContext, readsContext, featureContext, eval, evalName, null, stratLevelName, compVCs, evalSetBySample);
                    }
                }
            }

            if ( mergeEvals ) break; // stop processing the eval tracks
        }
    }

    private void processComp(ReferenceContext referenceContext, ReadsContext readsContext, FeatureContext featureContext, VariantContext eval, String evalName, FeatureInput<VariantContext> compInput, String stratLevelName, HashMap<FeatureInput<VariantContext>, HashMap<String, Collection<VariantContext>>> compVCs, Collection<VariantContext> evalSetBySample) {
        String compName = getNameForInput(compInput);

        // no sample stratification for comps
        final HashMap<String, Collection<VariantContext>> compSetHash = compInput == null ? null : compVCs.get(compInput);
        final Collection<VariantContext> compSet = (compSetHash == null || compSetHash.isEmpty()) ? Collections.<VariantContext>emptyList() : compVCs.get(compInput).values().iterator().next();

        // find the comp
        final VariantContext comp = findMatchingComp(eval, compSet);

        Collection<EvaluationContext> contextsForStratification;
        if (perFamilyIsEnabled)
            contextsForStratification = getEvaluationContexts(referenceContext, readsContext, featureContext, eval, evalName, comp, compName, null, stratLevelName);
        else {
            String familyID;
            if (stratLevelName.equals("all"))
                familyID = "all";
            else
                familyID = sampleDB.getSample(stratLevelName).getFamilyID();
            contextsForStratification = getEvaluationContexts(referenceContext, readsContext, featureContext, eval, evalName, comp, compName, stratLevelName, familyID);
        }
        for ( EvaluationContext nec : contextsForStratification ) {

            // eval against the comp
            synchronized (nec) {
                nec.apply(referenceContext, readsContext, featureContext, comp, eval);
            }

            // eval=null against all comps of different type that aren't bound to another eval
            for ( VariantContext otherComp : compSet ) {
                if ( otherComp != comp && ! compHasMatchingEval(otherComp, evalSetBySample) ) {
                    synchronized (nec) {
                        nec.apply(referenceContext, readsContext, featureContext, otherComp, null);
                    }
                }
            }
        }
    }

    /**
     * Given specific eval and comp VCs and the sample name, return an iterable
     * over all of the applicable state keys.
     *
     * this code isn't structured yet for efficiency.  Here we currently are
     * doing the following inefficient algorithm:
     *
     * for each strat:
     *   get list of relevant states that eval and comp according to strat
     *   add this list of states to a list of list states
     *
     * then
     *
     * ask the strat manager to look up all of the keys associated with the combinations
     * of these states.  For example, suppose we have a single variant S.  We have active
     * strats EvalFeatureInput, CompFeatureInput, and Novelty.  We produce a list that looks like:
     *
     *   L = [[Eval], [Comp], [All, Novel]]
     *
     * We then go through the strat manager tree to produce the keys associated with these states:
     *
     *   K = [0, 1] where EVAL x COMP x ALL = 0 and EVAL x COMP x NOVEL = 1
     *
     * It's clear that a better
     *
     *
     * @param referenceContext
     * @param readsContext
     * @param featureContext
     * @param eval
     * @param evalName
     * @param comp
     * @param compName
     * @param sampleName
     * @return
     */
    protected Collection<EvaluationContext> getEvaluationContexts(final ReferenceContext referenceContext,
                                                                  final ReadsContext readsContext,
                                                                  final FeatureContext featureContext,
                                                                  final VariantContext eval,
                                                                  final String evalName,
                                                                  final VariantContext comp,
                                                                  final String compName,
                                                                  final String sampleName,
                                                                  final String familyName) {
        final List<List<Object>> states = new LinkedList<List<Object>>();
        for ( final VariantStratifier vs : stratManager.getStratifiers() ) {
            states.add(vs.getRelevantStates(referenceContext, readsContext, featureContext, comp, compName, eval, evalName, sampleName, familyName));
        }
        return stratManager.values(states);
    }

    private boolean compHasMatchingEval(final VariantContext comp, final Collection<VariantContext> evals) {
        // find all of the matching comps
        for ( final VariantContext eval : evals ) {
            if ( eval != null && doEvalAndCompMatch(comp, eval, requireStrictAlleleMatch) != EvalCompMatchType.NO_MATCH )
                return true;
        }

        // nothing matched
        return false;
    }

    private enum EvalCompMatchType { NO_MATCH, STRICT, LENIENT }

    private EvalCompMatchType doEvalAndCompMatch(final VariantContext eval, final VariantContext comp, boolean requireStrictAlleleMatch) {
        if ( comp.getType() == VariantContext.Type.NO_VARIATION || eval.getType() == VariantContext.Type.NO_VARIATION )
            // if either of these are NO_VARIATION they are LENIENT matches
            return EvalCompMatchType.LENIENT;

        if ( comp.getType() != eval.getType() )
            return EvalCompMatchType.NO_MATCH;

        // find the comp which matches both the reference allele and alternate allele from eval
        final Allele altEval = eval.getAlternateAlleles().size() == 0 ? null : eval.getAlternateAllele(0);
        final Allele altComp = comp.getAlternateAlleles().size() == 0 ? null : comp.getAlternateAllele(0);
        if ((altEval == null && altComp == null) || (altEval != null && altEval.equals(altComp) && eval.getReference().equals(comp.getReference())))
            return EvalCompMatchType.STRICT;
        else
            return requireStrictAlleleMatch ? EvalCompMatchType.NO_MATCH : EvalCompMatchType.LENIENT;
    }

    private VariantContext findMatchingComp(final VariantContext eval, final Collection<VariantContext> comps) {
        // if no comps, return null
        if ( comps == null || comps.isEmpty() )
            return null;

        // if no eval, return any comp
        if ( eval == null )
            return comps.iterator().next();

        // find all of the matching comps
        VariantContext lenientMatch = null;
        for ( final VariantContext comp : comps ) {
            switch ( doEvalAndCompMatch(comp, eval, requireStrictAlleleMatch) ) {
                case STRICT:
                    return comp;
                case LENIENT:
                    if ( lenientMatch == null ) lenientMatch = comp;
                    break;
                case NO_MATCH:
                    // do nothing
            }
        }

        // nothing matched, just return lenientMatch, which might be null
        return lenientMatch;
    }

    @Override
    public Object onTraversalSuccess() {
        aggr.onComplete();

        logger.info("Finalizing variant report");

        // go through the evaluations and finalize them
        for ( final EvaluationContext nec : stratManager.values() )
            for ( final VariantEvaluator ve : nec.getVariantEvaluators() )
                ve.finalizeEvaluation();

        //send data to MetricsCollection
        CompOverlap compOverlap = null;
        IndelSummary indelSummary = null;
        CountVariants countVariants = null;
        MultiallelicSummary multiallelicSummary = null;
        TiTvVariantEvaluator tiTvVariantEvaluator = null;
        MetricsCollection metricsCollection = null;
        for(final EvaluationContext nec: stratManager.values()) {
            for(final VariantEvaluator ve : nec.getVariantEvaluators()) {
                if (ve instanceof CompOverlap)
                    compOverlap = (CompOverlap) ve;
                else if (ve instanceof IndelSummary)
                    indelSummary = (IndelSummary) ve;
                else if (ve instanceof CountVariants)
                    countVariants = (CountVariants) ve;
                else if (ve instanceof MultiallelicSummary)
                    multiallelicSummary = (MultiallelicSummary) ve;
                else if (ve instanceof TiTvVariantEvaluator)
                    tiTvVariantEvaluator = (TiTvVariantEvaluator) ve;
                else if (ve instanceof MetricsCollection)
                    metricsCollection = (MetricsCollection) ve;
            }

            if(metricsCollection != null)
                metricsCollection.setData(compOverlap.concordantRate, indelSummary.n_SNPs, countVariants.nSNPs, indelSummary.n_indels, multiallelicSummary.nIndels, indelSummary.insertion_to_deletion_ratio, countVariants.insertionDeletionRatio, tiTvVariantEvaluator.tiTvRatio);
        }

        try (PrintStream out = IOUtils.makePrintStreamMaybeGzipped(outFile)){
            VariantEvalReportWriter.writeReport(out, stratManager, stratManager.getStratifiers(), stratManager.get(0).getVariantEvaluators());
        }
        catch(IOException e) {
            throw new UserException.CouldNotCreateOutputFile(e.getMessage(), e);
        }

        return null;
    }

    // Accessors
    public Logger getLogger() { return logger; }

    public double getMinPhaseQuality() { return MIN_PHASE_QUALITY; }

    public int getSamplePloidy() { return ploidy; }
    public double getMendelianViolationQualThreshold() { return MENDELIAN_VIOLATION_QUAL_THRESHOLD; }

    public static String getAllSampleName() { return ALL_SAMPLE_NAME; }
    public static String getAllFamilyName() { return ALL_FAMILY_NAME; }

    public List<FeatureInput<VariantContext>> getKnowns() { return knowns; }

    public List<FeatureInput<VariantContext>> getEvals() { return evals; }

    public boolean isSubsettingToSpecificSamples() { return isSubsettingSamples; }
    public Set<String> getSampleNamesForEvaluation() { return sampleNamesForEvaluation; }

    public Set<String> getFamilyNamesForEvaluation() { return familyNamesForEvaluation; }

    public int getNumberOfSamplesForEvaluation() {
        if (sampleNamesForEvaluation!= null &&  !sampleNamesForEvaluation.isEmpty())
            return sampleNamesForEvaluation.size();
        else {
            return numSamplesFromArgument;
        }

    }
    public Set<String> getSampleNamesForStratification() { return sampleNamesForStratification; }

    public Set<String> getFamilyNamesForStratification() { return familyNamesForStratification; }

    public List<FeatureInput<VariantContext>> getComps() { return comps; }

    public Set<SortableJexlVCMatchExp> getJexlExpressions() { return jexlExpressions; }

    public StratifyingScale getAFScale() { return AFScale; }
    public boolean getCompAFStratifier() { return useCompAFStratifier; }

    public Set<String> getContigNames() {
        final TreeSet<String> contigs = new TreeSet<>();
        for( final SAMSequenceRecord r :  getSequenceDictionaryForDrivingVariants().getSequences()) {
            contigs.add(r.getSequenceName());
        }
        return contigs;
    }

    public boolean ignoreAC0Sites() {
        return ! keepSitesWithAC0;
    }

    /**
     * Entry-point function to initialize the samples database from input data
     */
    private SampleDB initializeSampleDB() {
        final SampleDBBuilder sampleDBBuilder = new SampleDBBuilder(PedigreeValidationType.STRICT);
        if (pedigreeFile != null)
            sampleDBBuilder.addSamplesFromPedigreeFiles(Collections.singletonList(pedigreeFile));

        Collection<String> samples = getSamplesForVariants();
        if (samples != null) {
            sampleDBBuilder.addSamplesFromSampleNames(samples);
        }

        return sampleDBBuilder.getFinalSampleDB();
    }

    public SampleDB getSampleDB() {
        return sampleDB;
    }

    /**
     * If an evaluator calls this method it must override {@link VariantEvaluator#requiresTerritoryToBeSpecified()} to return true.
     * @return either the size of the interval list given to the tool or the size of the reference given to the tool
     */
    public long getnProcessedLoci() {
        if(getTraversalIntervals() == null){
            throw new GATKException("BUG: One of the evaluators used should have overriden requiresTerritoryToBeSpecified, please report this to the developers." +
                    "\nEvaluators: " + stratManager.values()
                    .stream()
                    .flatMap(evaluator -> evaluator.getVariantEvaluators().stream())
                    .map(VariantEvaluator::getSimpleName)
                    .sorted()
                    .distinct()
                    .collect(Collectors.joining(", ")));
        }
        return getTraversalIntervals().stream().mapToLong(SimpleInterval::size).sum();
    }

    public FeatureInput<Feature> getKnownCNVsFile() {
        return knownCNVsFile;
    }

    @Override
    public List<? extends CommandLinePluginDescriptor<?>> getPluginDescriptors() {
        // Reason this override is needed:
        // Because getDefaultReadFilters() is nearly always included (see GATKTool), this pulls in SampleReadFilter,
        // which has an argument named "sample", which conflicts with the local argument.  we get an exception like:
        // org.broadinstitute.barclay.argparser.CommandLineException$CommandLineParserInternalException: [sample, sample] has already been used.
        // at org.broadinstitute.barclay.argparser.CommandLineArgumentParser.handleArgumentAnnotation(CommandLineArgumentParser.java:1002)
        // therefore override the method and return an empty list.
        return Collections.emptyList();
    }
}

{'repo_name': 'broadinstitute/gatk', 'stars': '843', 'repo_language': 'Java', 'file_name': 'testAFQCIntegration.expected.txt', 'mime_type': 'text/plain', 'hash': -3430645225376245461, 'source_dataset': 'data'}


// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

// Most of the code in this file comes from the default Roslyn Analyzer project template

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.Diagnostics;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Xunit;

namespace TestHelper
{
    /// <summary>
    /// Superclass of all Unit Tests for DiagnosticAnalyzers
    /// </summary>
    public abstract partial class DiagnosticVerifier
    {
        #region To be implemented by Test classes
        /// <summary>
        /// Get the CSharp analyzer being tested - to be implemented in non-abstract class
        /// </summary>
        protected virtual DiagnosticAnalyzer GetCSharpDiagnosticAnalyzer()
        {
            return null;
        }

        /// <summary>
        /// Get the Visual Basic analyzer being tested (C#) - to be implemented in non-abstract class
        /// </summary>
        protected virtual DiagnosticAnalyzer GetBasicDiagnosticAnalyzer()
        {
            return null;
        }
        #endregion

        #region Verifier wrappers

        /// <summary>
        /// Called to test a C# DiagnosticAnalyzer when applied on the single inputted string as a source
        /// Note: input a DiagnosticResult for each Diagnostic expected
        /// </summary>
        /// <param name="source">A class in the form of a string to run the analyzer on</param>
        /// <param name="expected"> DiagnosticResults that should appear after the analyzer is run on the source</param>
        protected void VerifyCSharpDiagnostic(string source, params DiagnosticResult[] expected)
        {
            VerifyDiagnostics(new[] { source }, LanguageNames.CSharp, GetCSharpDiagnosticAnalyzer(), expected);
        }

        /// <summary>
        /// Called to test a VB DiagnosticAnalyzer when applied on the single inputted string as a source
        /// Note: input a DiagnosticResult for each Diagnostic expected
        /// </summary>
        /// <param name="source">A class in the form of a string to run the analyzer on</param>
        /// <param name="expected">DiagnosticResults that should appear after the analyzer is run on the source</param>
        protected void VerifyBasicDiagnostic(string source, params DiagnosticResult[] expected)
        {
            VerifyDiagnostics(new[] { source }, LanguageNames.VisualBasic, GetBasicDiagnosticAnalyzer(), expected);
        }

        /// <summary>
        /// Called to test a C# DiagnosticAnalyzer when applied on the inputted strings as a source
        /// Note: input a DiagnosticResult for each Diagnostic expected
        /// </summary>
        /// <param name="sources">An array of strings to create source documents from to run the analyzers on</param>
        /// <param name="expected">DiagnosticResults that should appear after the analyzer is run on the sources</param>
        protected void VerifyCSharpDiagnostic(string[] sources, params DiagnosticResult[] expected)
        {
            VerifyDiagnostics(sources, LanguageNames.CSharp, GetCSharpDiagnosticAnalyzer(), expected);
        }

        /// <summary>
        /// Called to test a VB DiagnosticAnalyzer when applied on the inputted strings as a source
        /// Note: input a DiagnosticResult for each Diagnostic expected
        /// </summary>
        /// <param name="sources">An array of strings to create source documents from to run the analyzers on</param>
        /// <param name="expected">DiagnosticResults that should appear after the analyzer is run on the sources</param>
        protected void VerifyBasicDiagnostic(string[] sources, params DiagnosticResult[] expected)
        {
            VerifyDiagnostics(sources, LanguageNames.VisualBasic, GetBasicDiagnosticAnalyzer(), expected);
        }

        /// <summary>
        /// General method that gets a collection of actual diagnostics found in the source after the analyzer is run,
        /// then verifies each of them.
        /// </summary>
        /// <param name="sources">An array of strings to create source documents from to run the analyzers on</param>
        /// <param name="language">The language of the classes represented by the source strings</param>
        /// <param name="analyzer">The analyzer to be run on the source code</param>
        /// <param name="expected">DiagnosticResults that should appear after the analyzer is run on the sources</param>
        private void VerifyDiagnostics(string[] sources, string language, DiagnosticAnalyzer analyzer, params DiagnosticResult[] expected)
        {
            var diagnostics = GetSortedDiagnostics(sources, language, analyzer);
            VerifyDiagnosticResults(diagnostics, analyzer, expected);
        }

        #endregion

        #region Actual comparisons and verifications
        /// <summary>
        /// Checks each of the actual Diagnostics found and compares them with the corresponding DiagnosticResult in the array of expected results.
        /// Diagnostics are considered equal only if the DiagnosticResultLocation, Id, Severity, and Message of the DiagnosticResult match the actual diagnostic.
        /// </summary>
        /// <param name="actualResults">The Diagnostics found by the compiler after running the analyzer on the source code</param>
        /// <param name="analyzer">The analyzer that was being run on the sources</param>
        /// <param name="expectedResults">Diagnostic Results that should have appeared in the code</param>
        private static void VerifyDiagnosticResults(IEnumerable<Diagnostic> actualResults, DiagnosticAnalyzer analyzer, params DiagnosticResult[] expectedResults)
        {
            int expectedCount = expectedResults.Count();
            int actualCount = actualResults.Count();

            if (expectedCount != actualCount)
            {
                string diagnosticsOutput = actualResults.Any() ? FormatDiagnostics(analyzer, actualResults.ToArray()) : "    NONE.";

                Assert.True(false,
                    string.Format("Mismatch between number of diagnostics returned, expected \"{0}\" actual \"{1}\"\r\n\r\nDiagnostics:\r\n{2}\r\n", expectedCount, actualCount, diagnosticsOutput));
            }

            for (int i = 0; i < expectedResults.Length; i++)
            {
                var actual = actualResults.ElementAt(i);
                var expected = expectedResults[i];

                if (expected.Line == -1 && expected.Column == -1)
                {
                    if (actual.Location != Location.None)
                    {
                        Assert.True(false,
                            string.Format("Expected:\nA project diagnostic with No location\nActual:\n{0}",
                            FormatDiagnostics(analyzer, actual)));
                    }
                }
                else
                {
                    VerifyDiagnosticLocation(analyzer, actual, actual.Location, expected.Locations.First());
                    var additionalLocations = actual.AdditionalLocations.ToArray();

                    if (additionalLocations.Length != expected.Locations.Length - 1)
                    {
                        Assert.True(false,
                            string.Format("Expected {0} additional locations but got {1} for Diagnostic:\r\n    {2}\r\n",
                                expected.Locations.Length - 1, additionalLocations.Length,
                                FormatDiagnostics(analyzer, actual)));
                    }

                    for (int j = 0; j < additionalLocations.Length; ++j)
                    {
                        VerifyDiagnosticLocation(analyzer, actual, additionalLocations[j], expected.Locations[j + 1]);
                    }
                }

                if (actual.Id != expected.Id)
                {
                    Assert.True(false,
                        string.Format("Expected diagnostic id to be \"{0}\" was \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                            expected.Id, actual.Id, FormatDiagnostics(analyzer, actual)));
                }

                if (actual.Severity != expected.Severity)
                {
                    Assert.True(false,
                        string.Format("Expected diagnostic severity to be \"{0}\" was \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                            expected.Severity, actual.Severity, FormatDiagnostics(analyzer, actual)));
                }

                if (actual.GetMessage() != expected.Message)
                {
                    Assert.True(false,
                        string.Format("Expected diagnostic message to be \"{0}\" was \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                            expected.Message, actual.GetMessage(), FormatDiagnostics(analyzer, actual)));
                }
            }
        }

        /// <summary>
        /// Helper method to VerifyDiagnosticResult that checks the location of a diagnostic and compares it with the location in the expected DiagnosticResult.
        /// </summary>
        /// <param name="analyzer">The analyzer that was being run on the sources</param>
        /// <param name="diagnostic">The diagnostic that was found in the code</param>
        /// <param name="actual">The Location of the Diagnostic found in the code</param>
        /// <param name="expected">The DiagnosticResultLocation that should have been found</param>
        private static void VerifyDiagnosticLocation(DiagnosticAnalyzer analyzer, Diagnostic diagnostic, Location actual, DiagnosticResultLocation expected)
        {
            var actualSpan = actual.GetLineSpan();

            Assert.True(actualSpan.Path == expected.Path || (actualSpan.Path != null && actualSpan.Path.Contains("Test0.") && expected.Path.Contains("Test.")),
                string.Format("Expected diagnostic to be in file \"{0}\" was actually in file \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                    expected.Path, actualSpan.Path, FormatDiagnostics(analyzer, diagnostic)));

            var actualLinePosition = actualSpan.StartLinePosition;

            // Only check line position if there is an actual line in the real diagnostic
            if (actualLinePosition.Line > 0)
            {
                if (actualLinePosition.Line + 1 != expected.Line)
                {
                    Assert.True(false,
                        string.Format("Expected diagnostic to be on line \"{0}\" was actually on line \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                            expected.Line, actualLinePosition.Line + 1, FormatDiagnostics(analyzer, diagnostic)));
                }
            }

            // Only check column position if there is an actual column position in the real diagnostic
            if (actualLinePosition.Character > 0)
            {
                if (actualLinePosition.Character + 1 != expected.Column)
                {
                    Assert.True(false,
                        string.Format("Expected diagnostic to start at column \"{0}\" was actually at column \"{1}\"\r\n\r\nDiagnostic:\r\n    {2}\r\n",
                            expected.Column, actualLinePosition.Character + 1, FormatDiagnostics(analyzer, diagnostic)));
                }
            }
        }
        #endregion

        #region Formatting Diagnostics
        /// <summary>
        /// Helper method to format a Diagnostic into an easily readable string
        /// </summary>
        /// <param name="analyzer">The analyzer that this verifier tests</param>
        /// <param name="diagnostics">The Diagnostics to be formatted</param>
        /// <returns>The Diagnostics formatted as a string</returns>
        private static string FormatDiagnostics(DiagnosticAnalyzer analyzer, params Diagnostic[] diagnostics)
        {
            var builder = new StringBuilder();
            for (int i = 0; i < diagnostics.Length; ++i)
            {
                builder.AppendLine("// " + diagnostics[i].ToString());

                var analyzerType = analyzer.GetType();
                var rules = analyzer.SupportedDiagnostics;

                foreach (var rule in rules)
                {
                    if (rule != null && rule.Id == diagnostics[i].Id)
                    {
                        var location = diagnostics[i].Location;
                        if (location == Location.None)
                        {
                            builder.AppendFormat("GetGlobalResult({0}.{1})", analyzerType.Name, rule.Id);
                        }
                        else
                        {
                            Assert.True(location.IsInSource,
                                $"Test base does not currently handle diagnostics in metadata locations. Diagnostic in metadata: {diagnostics[i]}\r\n");

                            string resultMethodName = diagnostics[i].Location.SourceTree.FilePath.EndsWith(".cs") ? "GetCSharpResultAt" : "GetBasicResultAt";
                            var linePosition = diagnostics[i].Location.GetLineSpan().StartLinePosition;

                            builder.AppendFormat("{0}({1}, {2}, {3}.{4})",
                                resultMethodName,
                                linePosition.Line + 1,
                                linePosition.Character + 1,
                                analyzerType.Name,
                                rule.Id);
                        }

                        if (i != diagnostics.Length - 1)
                        {
                            builder.Append(',');
                        }

                        builder.AppendLine();
                        break;
                    }
                }
            }
            return builder.ToString();
        }
        #endregion
    }
}

{'repo_name': 'dotnet/aspnetcore', 'stars': '18362', 'repo_language': 'C#', 'file_name': 'Cpp.Common.targets', 'mime_type': 'text/plain', 'hash': -9136020910368254049, 'source_dataset': 'data'}


/*
 * Copyright (c) 2001, 2010, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 *
 */

#ifndef SHARE_VM_GC_IMPLEMENTATION_PARALLELSCAVENGE_PSMARKSWEEPDECORATOR_HPP
#define SHARE_VM_GC_IMPLEMENTATION_PARALLELSCAVENGE_PSMARKSWEEPDECORATOR_HPP

#include "gc_implementation/shared/mutableSpace.hpp"

//
// A PSMarkSweepDecorator is used to add "ParallelScavenge" style mark sweep operations
// to a MutableSpace.
//

class ObjectStartArray;

class PSMarkSweepDecorator: public CHeapObj<mtGC> {
 private:
  static PSMarkSweepDecorator* _destination_decorator;

 protected:
  MutableSpace* _space;
  ObjectStartArray* _start_array;
  HeapWord* _first_dead;
  HeapWord* _end_of_live;
  HeapWord* _compaction_top;
  size_t _allowed_dead_ratio;

  bool insert_deadspace(size_t& allowed_deadspace_words, HeapWord* q,
                        size_t word_len);

 public:
  PSMarkSweepDecorator(MutableSpace* space, ObjectStartArray* start_array,
                       size_t allowed_dead_ratio) :
    _space(space), _start_array(start_array),
    _allowed_dead_ratio(allowed_dead_ratio) { }

  // During a compacting collection, we need to collapse objects into
  // spaces in a given order. We want to fill space A, space B, and so
  // on. The code that controls that order is in the following methods.
  static void set_destination_decorator_tenured();
  static void set_destination_decorator_perm_gen();
  static void advance_destination_decorator();
  static PSMarkSweepDecorator* destination_decorator();

  // Accessors
  MutableSpace* space()                     { return _space; }
  ObjectStartArray* start_array()           { return _start_array; }

  HeapWord* compaction_top()                { return _compaction_top; }
  void set_compaction_top(HeapWord* value)  { _compaction_top = value; }

  size_t allowed_dead_ratio()               { return _allowed_dead_ratio; }
  void set_allowed_dead_ratio(size_t value) { _allowed_dead_ratio = value; }

  // Work methods
  void adjust_pointers();
  void precompact();
  void compact(bool mangle_free_space);
};

#endif // SHARE_VM_GC_IMPLEMENTATION_PARALLELSCAVENGE_PSMARKSWEEPDECORATOR_HPP

{'repo_name': 'codefollower/Open-Source-Research', 'stars': '197', 'repo_language': 'C++', 'file_name': 'SimpleExample.java', 'mime_type': 'text/x-java', 'hash': -7999416238364700114, 'source_dataset': 'data'}


package hashtag

import (
	"testing"

	"github.com/go-redis/redis/v8/internal/rand"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

func TestGinkgoSuite(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "hashtag")
}

var _ = Describe("CRC16", func() {
	// http://redis.io/topics/cluster-spec#keys-distribution-model
	It("should calculate CRC16", func() {
		tests := []struct {
			s string
			n uint16
		}{
			{"123456789", 0x31C3},
			{string([]byte{83, 153, 134, 118, 229, 214, 244, 75, 140, 37, 215, 215}), 21847},
		}

		for _, test := range tests {
			Expect(crc16sum(test.s)).To(Equal(test.n), "for %s", test.s)
		}
	})
})

var _ = Describe("HashSlot", func() {
	It("should calculate hash slots", func() {
		tests := []struct {
			key  string
			slot int
		}{
			{"123456789", 12739},
			{"{}foo", 9500},
			{"foo{}", 5542},
			{"foo{}{bar}", 8363},
			{"", 10503},
			{"", 5176},
			{string([]byte{83, 153, 134, 118, 229, 214, 244, 75, 140, 37, 215, 215}), 5463},
		}
		// Empty keys receive random slot.
		rand.Seed(100)

		for _, test := range tests {
			Expect(Slot(test.key)).To(Equal(test.slot), "for %s", test.key)
		}
	})

	It("should extract keys from tags", func() {
		tests := []struct {
			one, two string
		}{
			{"foo{bar}", "bar"},
			{"{foo}bar", "foo"},
			{"{user1000}.following", "{user1000}.followers"},
			{"foo{{bar}}zap", "{bar"},
			{"foo{bar}{zap}", "bar"},
		}

		for _, test := range tests {
			Expect(Slot(test.one)).To(Equal(Slot(test.two)), "for %s <-> %s", test.one, test.two)
		}
	})
})

{'repo_name': 'go-redis/redis', 'stars': '9241', 'repo_language': 'Go', 'file_name': 'redis.conf', 'mime_type': 'text/plain', 'hash': -7321893273651831520, 'source_dataset': 'data'}


glabel func_80B006B0
/* 02400 80B006B0 27BDFFE8 */  addiu   $sp, $sp, 0xFFE8           ## $sp = FFFFFFE8
/* 02404 80B006B4 AFBF0014 */  sw      $ra, 0x0014($sp)
/* 02408 80B006B8 AFA40018 */  sw      $a0, 0x0018($sp)
/* 0240C 80B006BC 0C2C0184 */  jal     func_80B00610
/* 02410 80B006C0 AFA5001C */  sw      $a1, 0x001C($sp)
/* 02414 80B006C4 0C2C0164 */  jal     func_80B00590
/* 02418 80B006C8 8FA40018 */  lw      $a0, 0x0018($sp)
/* 0241C 80B006CC 14400006 */  bne     $v0, $zero, .L80B006E8
/* 02420 80B006D0 8FA40018 */  lw      $a0, 0x0018($sp)
/* 02424 80B006D4 908E02D8 */  lbu     $t6, 0x02D8($a0)           ## 000002D8
/* 02428 80B006D8 55C00004 */  bnel    $t6, $zero, .L80B006EC
/* 0242C 80B006DC 8FBF0014 */  lw      $ra, 0x0014($sp)
/* 02430 80B006E0 0C2C0178 */  jal     func_80B005E0
/* 02434 80B006E4 00000000 */  nop
.L80B006E8:
/* 02438 80B006E8 8FBF0014 */  lw      $ra, 0x0014($sp)
.L80B006EC:
/* 0243C 80B006EC 27BD0018 */  addiu   $sp, $sp, 0x0018           ## $sp = 00000000
/* 02440 80B006F0 03E00008 */  jr      $ra
/* 02444 80B006F4 00000000 */  nop

{'repo_name': 'n64decomp/oot', 'stars': '265', 'repo_language': 'Assembly', 'file_name': 'func_808AAD04.s', 'mime_type': 'text/plain', 'hash': -3220898322871162609, 'source_dataset': 'data'}


#include <GuiMenu.au3>

_Main()

Func _Main()
	Local $hWnd, $hMain, $hFile

	; Open Notepad
	Run("notepad.exe")
	WinWaitActive("[CLASS:Notepad]")
	$hWnd = WinGetHandle("[CLASS:Notepad]")
	$hMain = _GUICtrlMenu_GetMenu($hWnd)
	$hFile = _GUICtrlMenu_GetItemSubMenu($hMain, 0)

	; Get/Set File menu data
	Writeln("File data: " & _GUICtrlMenu_GetMenuData($hFile))
	_GUICtrlMenu_SetMenuData($hFile, 1234)
	Writeln("File data: " & _GUICtrlMenu_GetMenuData($hFile))

EndFunc   ;==>_Main

; Write a line of text to Notepad
Func Writeln($sText)
	ControlSend("[CLASS:Notepad]", "", "Edit1", $sText & @CR)
EndFunc   ;==>Writeln

{'repo_name': 'forcedotcom/dataloader', 'stars': '269', 'repo_language': 'AutoIt', 'file_name': 'core.commands-3.6.1.v20120521-2329.pom', 'mime_type': 'text/xml', 'hash': 7049728112752923135, 'source_dataset': 'data'}


Summary: Baptist Health in Louisville, KY has notified 880 patients that some of their protected health information has potentially been accessed and stolen by hackers. The security breach was discovered on October 3, 2017, when irregular activity was detected on the email account of an employee. Baptist Health was able to determine that a third party sent a phishing email to the employee, who responded and disclosed login credentials allowing the email account to be accessed. Those login credentials were subsequently used by an unknown individual to gain access the email account. The email account contained the protected health information of 880 patients, although it is unclear whether any of the emails were viewed. The motive behind the attack may not have been to gain access to sensitive information.

{'repo_name': 'vz-risk/VCDB', 'stars': '346', 'repo_language': 'Jupyter Notebook', 'file_name': 'NOTE.txt', 'mime_type': 'text/plain', 'hash': 7064844509126182633, 'source_dataset': 'data'}


//===--- CapturePropagation.cpp - Propagate closure capture constants -----===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "capture-prop"

#include "polarphp/ast/GenericEnvironment.h"
#include "polarphp/demangling/Demangler.h"
#include "polarphp/pil/lang/PILCloner.h"
#include "polarphp/pil/lang/PILInstruction.h"
#include "polarphp/pil/lang/TypeSubstCloner.h"
#include "polarphp/pil/optimizer/analysis/ColdBlockInfo.h"
#include "polarphp/pil/optimizer/analysis/DominanceAnalysis.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/Generics.h"
#include "polarphp/pil/optimizer/utils/InstOptUtils.h"
#include "polarphp/pil/optimizer/utils/PILOptFunctionBuilder.h"
#include "polarphp/pil/optimizer/utils/SpecializationMangler.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/Debug.h"

using namespace polar;

STATISTIC(NumCapturesPropagated, "Number of constant captures propagated");

namespace {
/// Propagate constants through closure captures by specializing the partially
/// applied function.
/// Also optimize away partial_apply instructions where all partially applied
/// arguments are dead.
class CapturePropagation : public PILFunctionTransform
{
public:
   void run() override;

protected:
   bool optimizePartialApply(PartialApplyInst *PAI);
   PILFunction *specializeConstClosure(PartialApplyInst *PAI,
                                       PILFunction *SubstF);
   void rewritePartialApply(PartialApplyInst *PAI, PILFunction *SpecialF);
};
} // end anonymous namespace

static LiteralInst *getConstant(PILValue V) {
   if (auto I = dyn_cast<ThinToThickFunctionInst>(V))
      return getConstant(I->getOperand());
   if (auto I = dyn_cast<ConvertFunctionInst>(V))
      return getConstant(I->getOperand());
   return dyn_cast<LiteralInst>(V);
}

static bool isOptimizableConstant(PILValue V) {
   // We do not optimize string literals of length > 32 since we would need to
   // encode them into the symbol name for uniqueness.
   if (auto *SLI = dyn_cast<StringLiteralInst>(V))
      return SLI->getValue().size() <= 32;
   return true;
}

static bool isConstant(PILValue V) {
   V = getConstant(V);
   return V && isOptimizableConstant(V);
}

static std::string getClonedName(PartialApplyInst *PAI, IsSerialized_t Serialized,
                                 PILFunction *F) {
   auto P = demangling::SpecializationPass::CapturePropagation;
   mangle::FunctionSignatureSpecializationMangler Mangler(P, Serialized, F);

   // We know that all arguments are literal insts.
   unsigned argIdx = ApplySite(PAI).getCalleeArgIndexOfFirstAppliedArg();
   for (auto arg : PAI->getArguments()) {
      Mangler.setArgumentConstantProp(argIdx, getConstant(arg));
      ++argIdx;
   }
   return Mangler.mangle();
}

namespace {
/// Clone the partially applied function, replacing incoming arguments with
/// literal constants.
///
/// The cloned literals will retain the PILLocation from the partial apply's
/// caller, so the cloned function will have a mix of locations from different
/// functions.
class CapturePropagationCloner
   : public TypeSubstCloner<CapturePropagationCloner, PILOptFunctionBuilder> {
   using SuperTy =
   TypeSubstCloner<CapturePropagationCloner, PILOptFunctionBuilder>;
   friend class PILInstructionVisitor<CapturePropagationCloner>;
   friend class PILCloner<CapturePropagationCloner>;

   PILFunction *OrigF;
   bool IsCloningConstant;
public:
   CapturePropagationCloner(PILFunction *OrigF, PILFunction *NewF,
                            SubstitutionMap Subs)
      : SuperTy(*NewF, *OrigF, Subs), OrigF(OrigF), IsCloningConstant(false) {}

   void cloneClosure(OperandValueArrayRef Args);

protected:
   /// Literals cloned from the caller drop their location so the debug line
   /// tables don't senselessly jump around. As a placeholder give them the
   /// location of the newly cloned function.
   PILLocation remapLocation(PILLocation InLoc) {
      if (IsCloningConstant)
         return getBuilder().getFunction().getLocation();
      return InLoc;
   }

   /// Literals cloned from the caller take on the new function's debug scope.
   void postProcess(PILInstruction *Orig, PILInstruction *Cloned) {
      assert(IsCloningConstant == (Orig->getFunction() != OrigF) &&
             "Expect only cloned constants from the caller function.");
      PILClonerWithScopes<CapturePropagationCloner>::postProcess(Orig, Cloned);
   }

   const PILDebugScope *remapScope(const PILDebugScope *DS) {
      if (IsCloningConstant)
         return getBuilder().getFunction().getDebugScope();
      else
         return PILClonerWithScopes<CapturePropagationCloner>::remapScope(DS);
   }

   void cloneConstValue(PILValue Const);
};
} // end anonymous namespace

/// Clone a constant value. Recursively walk the operand chain through cast
/// instructions to ensure that all dependents are cloned. Note that the
/// original value may not belong to the same function as the one being cloned
/// by cloneClosure() (they may be from the partial apply caller).
void CapturePropagationCloner::cloneConstValue(PILValue Val) {
   assert(IsCloningConstant && "incorrect mode");

   if (isValueCloned(Val))
      return;

   // TODO: MultiValueInstruction?
   auto Inst = dyn_cast<SingleValueInstruction>(Val);
   if (!Inst)
      return;

   if (Inst->getNumOperands() > 0) {
      // Only handle single operands for simple recursion without a worklist.
      assert(Inst->getNumOperands() == 1 && "expected single-operand cast");
      cloneConstValue(Inst->getOperand(0));
   }
   visit(Inst);
}

/// Clone the original partially applied function into the new specialized
/// function, replacing some arguments with literals.
void CapturePropagationCloner::cloneClosure(
   OperandValueArrayRef PartialApplyArgs) {

   PILFunction &CloneF = getBuilder().getFunction();

   // Create the entry basic block with the function arguments.
   PILBasicBlock *OrigEntryBB = &*OrigF->begin();
   PILBasicBlock *ClonedEntryBB = CloneF.createBasicBlock();
   auto cloneConv = CloneF.getConventions();

   // Only clone the arguments that remain in the new function type. The trailing
   // arguments are now propagated through the partial apply.
   assert(!IsCloningConstant && "incorrect mode");

   SmallVector<PILValue, 4> entryArgs;
   entryArgs.reserve(OrigEntryBB->getArguments().size());

   unsigned ArgIdx = 0;
   for (unsigned NewArgEnd = cloneConv.getNumPILArguments(); ArgIdx != NewArgEnd;
        ++ArgIdx) {

      PILArgument *Arg = OrigEntryBB->getArgument(ArgIdx);

      PILValue MappedValue = ClonedEntryBB->createFunctionArgument(
         remapType(Arg->getType()), Arg->getDecl());
      entryArgs.push_back(MappedValue);
   }
   assert(OrigEntryBB->args_size() - ArgIdx == PartialApplyArgs.size()
          && "unexpected number of partial apply arguments");

   // Replace the rest of the old arguments with constants.
   getBuilder().setInsertionPoint(ClonedEntryBB);
   IsCloningConstant = true;
   for (PILValue PartialApplyArg : PartialApplyArgs) {
      assert(isConstant(PartialApplyArg) &&
             "expected a constant arg to partial apply");

      cloneConstValue(PartialApplyArg);

      // The PartialApplyArg from the caller is now mapped to its cloned
      // instruction.  Also map the original argument to the cloned instruction.
      entryArgs.push_back(getMappedValue(PartialApplyArg));
      ++ArgIdx;
   }
   IsCloningConstant = false;

   // Clear information about cloned values from the caller function.
   clearClonerState();

   // Visit original BBs in depth-first preorder, starting with the
   // entry block, cloning all instructions and terminators.
   cloneFunctionBody(OrigF, ClonedEntryBB, entryArgs);
}

CanPILFunctionType getPartialApplyInterfaceResultType(PartialApplyInst *PAI) {
   // The new partial_apply will no longer take any arguments--they are all
   // expressed as literals. So its callee signature will be the same as its
   // return signature.
   auto FTy = PAI->getType().castTo<PILFunctionType>();
   assert(!PAI->hasSubstitutions() ||
          !PAI->getSubstitutionMap().hasArchetypes());
   FTy = cast<PILFunctionType>(
      FTy->mapTypeOutOfContext()->getCanonicalType());
   auto NewFTy = FTy;
   return NewFTy;
}

/// Given a partial_apply instruction, create a specialized callee by removing
/// all constant arguments and adding constant literals to the specialized
/// function body.
PILFunction *CapturePropagation::specializeConstClosure(PartialApplyInst *PAI,
                                                        PILFunction *OrigF) {
   IsSerialized_t Serialized = IsNotSerialized;
   if (PAI->getFunction()->isSerialized() && OrigF->isSerialized())
      Serialized = IsSerializable;

   std::string Name = getClonedName(PAI, Serialized, OrigF);

   // See if we already have a version of this function in the module. If so,
   // just return it.
   if (auto *NewF = OrigF->getModule().lookUpFunction(Name)) {
      assert(NewF->isSerialized() == Serialized);
      LLVM_DEBUG(llvm::dbgs()
                    << "  Found an already specialized version of the callee: ";
                    NewF->printName(llvm::dbgs()); llvm::dbgs() << "\n");
      return NewF;
   }

   // The new partial_apply will no longer take any arguments--they are all
   // expressed as literals. So its callee signature will be the same as its
   // return signature.
   auto NewFTy = getPartialApplyInterfaceResultType(PAI);
   NewFTy = NewFTy->getWithRepresentation(PILFunctionType::Representation::Thin);

   GenericEnvironment *GenericEnv = nullptr;
   if (NewFTy->getInvocationGenericSignature())
      GenericEnv = OrigF->getGenericEnvironment();
   PILOptFunctionBuilder FuncBuilder(*this);
   PILFunction *NewF = FuncBuilder.createFunction(
      PILLinkage::Shared, Name, NewFTy, GenericEnv, OrigF->getLocation(),
      OrigF->isBare(), OrigF->isTransparent(), Serialized, IsNotDynamic,
      OrigF->getEntryCount(), OrigF->isThunk(), OrigF->getClassSubclassScope(),
      OrigF->getInlineStrategy(), OrigF->getEffectsKind(),
      /*InsertBefore*/ OrigF, OrigF->getDebugScope());
   if (!OrigF->hasOwnership()) {
      NewF->setOwnershipEliminated();
   }
   LLVM_DEBUG(llvm::dbgs() << "  Specialize callee as ";
                 NewF->printName(llvm::dbgs());
                 llvm::dbgs() << " " << NewFTy << "\n");

   LLVM_DEBUG(if (PAI->hasSubstitutions()) {
      llvm::dbgs() << "CapturePropagation of generic partial_apply:\n";
      PAI->dumpInContext();
   });
   CapturePropagationCloner cloner(OrigF, NewF, PAI->getSubstitutionMap());
   cloner.cloneClosure(PAI->getArguments());
   assert(OrigF->getDebugScope()->Parent != NewF->getDebugScope()->Parent);
   return NewF;
}

void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
                                             PILFunction *SpecialF) {
   LLVM_DEBUG(llvm::dbgs() << "\n  Rewriting a partial apply:\n";
                 OrigPAI->dumpInContext();
                 llvm::dbgs() << "   with special function: "
                              << SpecialF->getName() << "\n";
                 llvm::dbgs() << "\nThe function being rewritten is:\n";
                 OrigPAI->getFunction()->dump());

   PILBuilderWithScope Builder(OrigPAI);
   auto FuncRef = Builder.createFunctionRef(OrigPAI->getLoc(), SpecialF);
   auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(), FuncRef,
                                                  OrigPAI->getType());
   OrigPAI->replaceAllUsesWith(T2TF);
   // Remove any dealloc_stack users.
   SmallVector<Operand*, 16> Uses(T2TF->getUses());
   for (auto *Use : Uses)
      if (auto *DS = dyn_cast<DeallocStackInst>(Use->getUser()))
         DS->eraseFromParent();
   recursivelyDeleteTriviallyDeadInstructions(OrigPAI, true);
   LLVM_DEBUG(llvm::dbgs() << "  Rewrote caller:\n" << *T2TF);
}

/// For now, we conservative only specialize if doing so can eliminate dynamic
/// dispatch.
///
/// TODO: Check for other profitable constant propagation, like builtin compare.
static bool isProfitable(PILFunction *Callee) {
   PILBasicBlock *EntryBB = &*Callee->begin();
   for (auto *Arg : EntryBB->getArguments()) {
      for (auto *Operand : Arg->getUses()) {
         if (FullApplySite FAS = FullApplySite::isa(Operand->getUser())) {
            if (FAS.getCallee() == Operand->get())
               return true;
         }
      }
   }
   return false;
}

/// Returns true if block \p BB only contains a return or throw of the first
/// block argument and side-effect-free instructions.
static bool onlyContainsReturnOrThrowOfArg(PILBasicBlock *BB) {
   for (PILInstruction &I : *BB) {
      if (isa<ReturnInst>(&I) || isa<ThrowInst>(&I)) {
         PILValue RetVal = I.getOperand(0);
         return BB->getNumArguments() == 1 && RetVal == BB->getArgument(0);
      }
      if (I.mayHaveSideEffects() || isa<TermInst>(&I))
         return false;
   }
   llvm_unreachable("should have seen a terminator instruction");
}

/// Checks if \p Orig is a thunk which calls another function but without
/// passing the trailing \p numDeadParams dead parameters.
/// If a generic specialization was performed for a generic capture,
/// GenericSpecialized contains a tuple:
/// (new specialized function, old function)
static PILFunction *getSpecializedWithDeadParams(
   PILOptFunctionBuilder &FuncBuilder,
   PartialApplyInst *PAI, PILFunction *Orig, int numDeadParams,
   std::pair<PILFunction *, PILFunction *> &GenericSpecialized) {
   PILBasicBlock &EntryBB = *Orig->begin();
   unsigned NumArgs = EntryBB.getNumArguments();

   // Check if all dead parameters have trivial types. We don't support non-
   // trivial types because it's very hard to find places where we can release
   // those parameters (as a replacement for the removed partial_apply).
   // TODO: maybe we can skip this restriction when we have semantic ARC.
   for (unsigned Idx = NumArgs - numDeadParams; Idx < NumArgs; ++Idx) {
      PILType ArgTy = EntryBB.getArgument(Idx)->getType();
      if (!ArgTy.isTrivial(*Orig))
         return nullptr;
   }
   PILFunction *Specialized = nullptr;
   PILValue RetValue;

   // Check all instruction of the entry block.
   for (PILInstruction &I : EntryBB) {
      if (auto FAS = FullApplySite::isa(&I)) {
         // Check if this is the call of the specialized function.
         // If the original partial_apply didn't have substitutions,
         // also the specialized function must be not generic.
         if (!PAI->hasSubstitutions() && FAS.hasSubstitutions())
            return nullptr;

         // Is it the only call?
         if (Specialized)
            return nullptr;

         Specialized = FAS.getReferencedFunctionOrNull();
         if (!Specialized)
            return nullptr;

         // Check if parameters are passes 1-to-1
         unsigned NumArgs = FAS.getNumArguments();
         if (EntryBB.getNumArguments() - numDeadParams != NumArgs)
            return nullptr;

         for (unsigned Idx = 0; Idx < NumArgs; ++Idx) {
            if (FAS.getArgument(Idx) != (ValueBase *)EntryBB.getArgument(Idx))
               return nullptr;
         }

         if (auto *TAI = dyn_cast<TryApplyInst>(&I)) {
            // Check the normal and throw blocks of the try_apply.
            if (onlyContainsReturnOrThrowOfArg(TAI->getNormalBB()) &&
                onlyContainsReturnOrThrowOfArg(TAI->getErrorBB()))
               return Specialized;
            return nullptr;
         }
         assert(isa<ApplyInst>(&I) && "unknown FullApplySite instruction");
         RetValue = cast<ApplyInst>(&I);
         continue;
      }
      if (auto *RI = dyn_cast<ReturnInst>(&I)) {
         // Check if we return the result of the apply.
         if (RI->getOperand() != RetValue)
            return nullptr;
         continue;
      }
      if (I.mayHaveSideEffects() || isa<TermInst>(&I))
         return nullptr;
   }

   auto Rep = Specialized->getLoweredFunctionType()->getRepresentation();
   if (getPILFunctionLanguage(Rep) != PILFunctionLanguage::Swift)
      return nullptr;

   GenericSpecialized = std::make_pair(nullptr, nullptr);

   if (PAI->hasSubstitutions()) {
      if (Specialized->isExternalDeclaration())
         return nullptr;
      if (!Orig->shouldOptimize())
         return nullptr;

      // Perform a generic specialization of the Specialized function.
      ReabstractionInfo ReInfo(
         FuncBuilder.getModule().getTypePHPModule(),
         FuncBuilder.getModule().isWholeModule(), ApplySite(), Specialized,
         PAI->getSubstitutionMap(), Specialized->isSerialized(),
         /* ConvertIndirectToDirect */ false);
      GenericFuncSpecializer FuncSpecializer(FuncBuilder,
                                             Specialized,
                                             ReInfo.getClonerParamSubstitutionMap(),
                                             ReInfo);

      PILFunction *GenericSpecializedFunc = FuncSpecializer.trySpecialization();
      if (!GenericSpecializedFunc)
         return nullptr;
      GenericSpecialized = std::make_pair(GenericSpecializedFunc, Specialized);
      return GenericSpecializedFunc;
   }
   return Specialized;
}

bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   PILFunction *SubstF = PAI->getReferencedFunctionOrNull();
   if (!SubstF)
      return false;
   if (SubstF->isExternalDeclaration())
      return false;

   if (PAI->hasSubstitutions() && PAI->getSubstitutionMap().hasArchetypes()) {
      LLVM_DEBUG(llvm::dbgs()
                    << "CapturePropagation: cannot handle partial specialization "
                       "of partial_apply:\n";
                    PAI->dumpInContext());
      return false;
   }

   // First possibility: Is it a partial_apply where all partially applied
   // arguments are dead?
   std::pair<PILFunction *, PILFunction *> GenericSpecialized;
   PILOptFunctionBuilder FuncBuilder(*this);
   if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder,
                                                    PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) {
      rewritePartialApply(PAI, NewFunc);
      if (GenericSpecialized.first) {
         // Notify the pass manager about the new function.
         addFunctionToPassManagerWorklist(GenericSpecialized.first,
                                          GenericSpecialized.second);
      }
      return true;
   }

   // Second possibility: Are all partially applied arguments constant?
   for (auto Arg : PAI->getArguments()) {
      if (!isConstant(Arg))
         return false;
   }
   if (!isProfitable(SubstF))
      return false;

   LLVM_DEBUG(llvm::dbgs() << "Specializing closure for constant arguments:\n"
                           << "  " << SubstF->getName() << "\n"
                           << *PAI);
   ++NumCapturesPropagated;
   PILFunction *NewF = specializeConstClosure(PAI, SubstF);
   rewritePartialApply(PAI, NewF);

   addFunctionToPassManagerWorklist(NewF, SubstF);
   return true;
}

void CapturePropagation::run() {
   DominanceAnalysis *DA = PM->getAnalysis<DominanceAnalysis>();
   auto *F = getFunction();
   bool HasChanged = false;

   // Don't optimize functions that are marked with the opt.never attribute.
   if (!F->shouldOptimize())
      return;

   // Cache cold blocks per function.
   ColdBlockInfo ColdBlocks(DA);
   for (auto &BB : *F) {
      if (ColdBlocks.isCold(&BB))
         continue;

      auto I = BB.begin();
      while (I != BB.end()) {
         PILInstruction *Inst = &*I;
         ++I;
         if (auto *PAI = dyn_cast<PartialApplyInst>(Inst))
            HasChanged |= optimizePartialApply(PAI);
      }
   }
   if (HasChanged) {
      invalidateAnalysis(PILAnalysis::InvalidationKind::Everything);
   }
}

PILTransform *polar::createCapturePropagation() {
   return new CapturePropagation();
}

{'repo_name': 'polarphp/polarphp', 'stars': '1040', 'repo_language': 'C++', 'file_name': 'Commands.cpp', 'mime_type': 'text/x-c++', 'hash': 2042897897912194677, 'source_dataset': 'data'}