package edu.utexas.cs.utopia.cortado.ccrs;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterators;
import edu.utexas.cs.utopia.cortado.staticanalysis.rwsetanalysis.MayRWSetAnalysis;
import edu.utexas.cs.utopia.cortado.staticanalysis.rwsetanalysis.MemoryLocations;
import edu.utexas.cs.utopia.cortado.staticanalysis.singletons.CachedMayRWSetAnalysis;
import edu.utexas.cs.utopia.cortado.util.graph.SootGraphConverter;
import edu.utexas.cs.utopia.cortado.util.graph.StronglyConnectedComponents;
import edu.utexas.cs.utopia.cortado.util.soot.ExceptionalUnitGraphCache;
import edu.utexas.cs.utopia.cortado.util.soot.SootUnitUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import soot.Body;
import soot.SootClass;
import soot.SootMethod;
import soot.Unit;
import soot.jimple.IdentityStmt;
import soot.toolkits.graph.Block;
import soot.toolkits.graph.CompleteBlockGraph;
import soot.toolkits.graph.ExceptionalUnitGraph;
import soot.toolkits.graph.MHGDominatorsFinder;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import static edu.utexas.cs.utopia.cortado.staticanalysis.rwsetanalysis.MayRWSetAnalysis.getMemoryLocationsForFields;
import static edu.utexas.cs.utopia.cortado.staticanalysis.rwsetanalysis.MayRWSetAnalysis.getThisMemoryLocs;

public class PointerAnalysisFragmentPartitioner extends FragmentPartitioner
{
    private static final Logger log = LoggerFactory.getLogger(PointerAnalysisFragmentPartitioner.class);

    private final MayRWSetAnalysis mayRWSetAnalysis = CachedMayRWSetAnalysis.getInstance().getMayRWSetAnalysis();

    private final Set<MemoryLocations> mtrFldMemLocs;

    private final Set<MemoryLocations> mtrThisMemLocs;

    public PointerAnalysisFragmentPartitioner(SootClass mtr)
    {
        this.mtrFldMemLocs = getMemoryLocationsForFields(mtr);
        this.mtrThisMemLocs = getThisMemoryLocs(mtr);
    }

    @Override
    List<Fragment> partition(CCR ccr)
    {
        SootMethod ccrMeth = ccr.getAtomicSection();
        Body methBody = ccrMeth.getActiveBody();

        if (!methBody.getTraps().isEmpty())
            log.warn("Ignoring traps inside ccr: " + ccrMeth);

        CompleteBlockGraph blockGraph = new CompleteBlockGraph(methBody);
        StronglyConnectedComponents<Block> sccs = new StronglyConnectedComponents<>(SootGraphConverter.convertToGuavaGraph(blockGraph));
        final ExceptionalUnitGraph cfg = ExceptionalUnitGraphCache.getInstance().getOrCompute(ccr.getAtomicSection());

        MHGDominatorsFinder<Block> dominators = new MHGDominatorsFinder<>(blockGraph);

        Set<Block> loopHeads = new HashSet<>();
        blockGraph.forEach(b -> {
            List<Block> bDoms = dominators.getDominators(b);
            b.getSuccs()
             .forEach(succ -> {
                 if (bDoms.contains(succ))
                     loopHeads.add(succ);
             });
        });

        // go ahead and figure out the exit units
        Set<Unit> explicitExitUnits = SootUnitUtils.getExplicitExitUnits(ccr.getAtomicSection());
        List<Fragment> fragPartition = new ArrayList<>();
        final Set<Unit> waituntilFragUnits = new HashSet<>();
        if(ccr.hasGuard())
        {
            waituntilFragUnits.addAll(ccr.getWaituntilFragment().getAllUnits());
        }

        ImmutableList<ImmutableList<Block>> sccsComp = sccs.getComponents();
        assert !sccs.getComponents().isEmpty();
        for (int i = sccsComp.size() - 1; i >= 0; i--)
        {
            ImmutableList<Block> scc = sccsComp.get(i);
            final boolean sccMeetsWaituntilFrag = scc.parallelStream()
                    .map(Iterable::spliterator)
                    .flatMap(it -> StreamSupport.stream(it, true))
                    .anyMatch(waituntilFragUnits::contains);
            if(sccMeetsWaituntilFrag)
            {
                final boolean sccContainedInWaituntilFrag = scc.stream()
                        .map(Iterable::spliterator)
                        .flatMap(it -> StreamSupport.stream(it, false))
                        .collect(Collectors.collectingAndThen(Collectors.toList(),
                                 waituntilFragUnits::containsAll));
                if(!sccContainedInWaituntilFrag)
                {
                    throw new IllegalStateException("SCC meets, but is not contained by, waituntil fragment");
                }
                continue;
            }

            // Group loops together (except for the while(!guard) yield() loops generated by the CCR)
            if (!Collections.disjoint(loopHeads, scc))
            {
                // Bundle all blocks inside loops in the same SCC.
                List<Block> entryBlocks = scc.stream()
                                             .filter(b -> !scc.containsAll(b.getPreds()))
                                             .collect(Collectors.toList());

                List<Block> exitBlocks = scc.stream()
                                            .filter(b -> !scc.containsAll(b.getSuccs()))
                                            .collect(Collectors.toList());

                assert entryBlocks.size() == 1;

                Unit fragStart = entryBlocks.get(0).getHead();
                Set<Unit> fragExits = exitBlocks.stream()
                        .map(Block::getTail)
                        .collect(Collectors.toSet());

                final boolean sccYields = scc.stream()
                        .map(Block::iterator)
                        .anyMatch(it -> Iterators.any(it, SootUnitUtils::unitIsWaiting));

                // case that we're yielding on a guard
                if(sccYields)
                {
                    scc.stream()
                            .map(Block::spliterator)
                            .map(unitsInBlockIter -> StreamSupport.stream(unitsInBlockIter, false))
                            .map(unitsInBlockStream -> unitsInBlockStream.collect(Collectors.toList()))
                            .map(unitsInBlock -> new Fragment(ccr, unitsInBlock))
                            .forEach(fragPartition::add);
                }
                else
                {
                    // replace units which exit the method by their predecessors inside the scc
                    fragExits.stream()
                            .filter(explicitExitUnits::contains)
                            .map(cfg::getPredsOf)
                            .filter(ut -> scc.stream().anyMatch(b -> Iterators.contains(b.iterator(), ut)))
                            .flatMap(Collection::stream)
                            .forEach(fragExits::add);
                    fragExits.removeIf(explicitExitUnits::contains);

                    fragPartition.add(new Fragment(ccr, fragStart, fragExits));
                }
            }
            else
            {
                Block b = scc.get(0);
                Unit tail = b.getTail();

                Unit fragStart = null;
                Unit lastFragEnd = null;
                for (Unit u : b)
                {
                    if (u instanceof IdentityStmt) continue;

                    if (u == tail)
                    {
                        if (lastFragEnd != u && !explicitExitUnits.contains(u))
                        {
                            fragPartition.add(new Fragment(ccr, tail));
                        }
                        break;
                    }


                    if (fragStart == null)
                        fragStart = u;

                    Unit succOfU = b.getSuccOf(u);
                    Collection<MemoryLocations> uWriteSet = mayRWSetAnalysis.writeSet(ccr.getAtomicSection(), u);
                    if (succOfU == tail ||
                        MemoryLocations.mayIntersectAny(uWriteSet, mtrFldMemLocs) ||
                        MemoryLocations.mayIntersectAny(uWriteSet, mtrThisMemLocs))
                    {
                        Unit fragEnd = succOfU == tail && !explicitExitUnits.contains(tail) ? tail: u;
                        fragPartition.add(new Fragment(ccr, fragStart, Collections.singleton(fragEnd)));
                        lastFragEnd = fragEnd;
                        fragStart = null;
                    }
                }
            }
        }

        return fragPartition;
    }
}
