/*
 * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *   - Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *
 *   - Redistributions in binary form must reproduce the above copyright
 *     notice, this list of conditions and the following disclaimer in the
 *     documentation and/or other materials provided with the distribution.
 *
 *   - Neither the name of Oracle nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

import java.nio.file.*;
import java.util.*;
import java.util.stream.Stream;
import java.util.regex.*;

// This script takes traces produced by TraceEscapeAnalysis and outputs the escaping allocations in a more readable format
//
// To use:
// 1. Generate a trace file by running with: "-XX:CompileCommand=TraceEscapeAnalysis,package.<class>::<method>" (need fastdebug build)
// 2. Pipe output to a file. Remove any content that does no belong to the EA trace.
// 3. Run this parser: java TraceEAParser.java ./path/to/trace/file
//
// You might encounter errors for malformed trace files, such as nodes not being found. In that case there is
// usually a newline missing from the log (e.g. due to other JMH output). Search for the node and then add the newline
// manually.
//
// Output should look something like:
//
// Escaping allocations:
//
// JavaObject(17) allocation in: ConfinedSession::<init> @ bci:2 (line 55)
//   -> Field(35)
//   -> JavaObject(16)
//   -> Field(39)
//   -> JavaObject(18)
//   -> LocalVar(95)
//   -> LocalVar(139)
//   Reason: EscapesAsArg[callNode= 343  CallStaticJava  === 500 501 502 48 1 (503 504 491 100 101 1 1 1 102 1 1 103 212 1 105 1 102 1 503 1 1 1 1 ) [[ 472 179 278 280 ]] # Static  java.lang.foreign.SegmentAllocator::copyArrayWithSwapIfNeeded java/lang/Object * ( java/lang/Object:NotNull *, java/lang/Object *, java/lang/Object * ) SegmentAllocator::allocateFrom @ bci:3 (line 268) AllocFromTest::alloc_confined @ bci:12 (line 72) AllocFromTest_alloc_confined_jmhTest::alloc_confined_avgt_jmhStub @ bci:17 (line 186) !jvms: SegmentAllocator::allocateFrom @ bci:3 (line 268) AllocFromTest::alloc_confined @ bci:12 (line 72) AllocFromTest_alloc_confined_jmhTest::alloc_confined_avgt_jmhStub @ bci:17 (line 186)]
//
// JavaObject(26) allocation in: NativeMemorySegmentImpl::makeNativeSegment @ bci:79 (line 132)
//   -> LocalVar(85)
//   -> LocalVar(117)
//   -> LocalVar(72)
//   Reason: PropagatedFromNode[node=LocalVar(72) [ [ ]]      9  Return  === 54 55 56 47 57 returns 58  [[ 0 ]]]
//
// JavaObject(31) allocation in: NativeMemorySegmentImpl::makeNativeSegment @ bci:112 (line 136)
//   Reason: MergedWithObject[other=JavaObject(1) [ [ ]]    230  ConP  === 0  [[ 521 133 521 699 140 135 453 455 134 139 ]]  #null]
//
// JavaObject(20) allocation in: ConfinedSession::<init> @ bci:2 (line 55)
//   -> Field(66)
//   -> JavaObject(19)
//   -> Field(45)
//   -> JavaObject(22)
//   -> LocalVar(107)
//   -> LocalVar(151)
//   Reason: BlackHole[]
//
public class TraceEAParser {

    private static final String ESCAPE_TYPE_PAT_NC = "(?:NoEscape|ArgEscape|GlobalEscape)";
    private static final String STATE_PAT_NC = ESCAPE_TYPE_PAT_NC + "\\(" + ESCAPE_TYPE_PAT_NC + "\\)";
    private static final String NODE_PAT = "(?<nodeType>LocalVar|JavaObject|Field)\\((?<nodeIdx>\\d+)\\) " + STATE_PAT_NC + "(?: NSR)?";

    private static final Pattern PROP_PAT = Pattern.compile(
        NODE_PAT
        + " -> " + STATE_PAT_NC
        + " propagated from: (?:LocalVar|JavaObject|Field)\\((?<srcNodeIdx>\\d+)\\)");

    private static final Pattern ESCAPE_AS_ARG_PAT = Pattern.compile(
        NODE_PAT
        + " -> " + STATE_PAT_NC
        + " escapes as arg to: (?<escapeNode>.*)");

    private static final Pattern MERGED_PAT = Pattern.compile(
        NODE_PAT
        + " is NSR. is merged with other object: "
        + "(?:LocalVar|JavaObject|Field)\\((?<srcNodeIdx>\\d+)\\)(?: NSR)?");

    private static final Pattern UNKNOWN_FIELD_OFFSET_PAT = Pattern.compile(
        NODE_PAT
        + " is NSR. has field with unknown offset");

    private static final Pattern BLACKHOLE_PAT = Pattern.compile(
        NODE_PAT
        + " -> " + STATE_PAT_NC
        + " blackhole");

    private static final Pattern INIT_PAT = Pattern.compile(
        NODE_PAT
        + " (?<nodeDump>.*)");

    private static final Pattern ALLOC_SOURCE_PAT = Pattern.compile(
        "([^ ]+ @ bci:\\d+ \\(line \\d+\\))");

    private static final Pattern FIRST_CALL_PAT = Pattern.compile(
        ".*#\s+(Static)?\s+(?<callname>.+?\s+\\(line \\d+\\))");

    private static final String INV_START_LINE = "JavaObject(0) GlobalEscape(GlobalEscape) NSR is NSR. Phantom object";

    private static void fatal(String message) {
        System.err.println("ERROR: " + message);
        System.exit(-1);
    }

    public static void main(String[] args) throws Throwable {
        String fileName = null;
        boolean debug = false;
        for (int i = 0; i < args.length; i++) {
            switch (args[i]) {
                case "-d" -> { debug = true; }
                case "-f" -> { fileName = args[++i]; }
                default -> { fileName = args[i]; }
            }
        }
        if (fileName == null) {
            fatal("Needs file to parse. " );
        }

        Path filePath = Path.of(fileName);
        if (!Files.exists(filePath)) {
            fatal("File does not exist: " + filePath);
        }

        List<EAInvocation> invocations = new ArrayList<>();
        try (Stream<String> lines = Files.lines(filePath)) {
            Iterator<String> cursor = lines.iterator();
            while (cursor.hasNext()) {
                invocations.add(parseInv(cursor));
            }
        } catch (ParseError pe) {
            if (debug) {
                pe.printStackTrace();
            }
            fatal(pe.getMessage());
        }

        EAInvocation last = invocations.get(invocations.size() - 1);

        System.out.println("Escaping allocations:");
        System.out.println();
        for (Node alloc : last.allocations()) {
            Reason reason = last.propagations().get(alloc);
            if (reason == null) {
                continue; // was scalarized
            }

            System.out.print(alloc.header() + " allocation in: ");
            Matcher m = ALLOC_SOURCE_PAT.matcher(alloc.fullDump());
            if (m.find()) {
                System.out.println(m.group(1));
            } else {
                System.out.println("unknown");
            }

            Reason lastProp = null;
            while (reason != null) {
                lastProp = reason;
                if (reason instanceof Reason.PropagatedFromNode fn) {
                    System.out.println("  -> " + fn.node().header());
                    reason = last.propagations().get(fn.node());
                } else {
                    break;
                }
            }
            System.out.println("  Reason: " + lastProp.simpleReason());

            System.out.println();
            System.out.println();
        }
    }

    private static EAInvocation parseInv(Iterator<String> cursor) {
        List<Node> allocations = new ArrayList<>();
        Map<Node, Reason> propagations = new HashMap<>();
        Map<Integer, Node> idxToNode = new HashMap<>();

        idxToNode.put(0, new Node(NodeType.JavaObject, 0, "Phantom object"));

        String line;
        while (!(line = cursor.next()).startsWith("+++++ Initial worklist")) {
            // perform some basic validation
            if(!line.matches("^JavaObject\\(\\d+\\).*$")) {
                throw new ParseError("Unrecognized line: " + line);
            }
        }
        
        while (!(line = cursor.next()).startsWith("+++++ Calculating escape states")) {
            Matcher m = INIT_PAT.matcher(line);
            if (!m.find()) {
                throw new ParseError("Can not parse line: " + line);
            }

            NodeType nodeType = NodeType.valueOf(m.group("nodeType"));
            int nodeIdx = Integer.parseInt(m.group("nodeIdx"));
            String nodeDump = m.group("nodeDump");
            Node node = new Node(nodeType, nodeIdx, nodeDump);
            idxToNode.put(nodeIdx, node);

            if (nodeType == NodeType.JavaObject && nodeDump.contains("Allocate  ===")) {
                allocations.add(node);
            }
        }

        while (cursor.hasNext() && !(line = cursor.next()).startsWith(INV_START_LINE)) {
            if (line.contains("propagated from")) {
                Matcher m = parseLine(PROP_PAT, line);

                Node node = getNode(m.group("nodeIdx"), idxToNode, line);
                Node srcNode = getNode(m.group("srcNodeIdx"), idxToNode, line);

                propagations.put(node, new Reason.PropagatedFromNode(srcNode));
            } else if (line.contains("escapes as arg to")) {
                Matcher m = parseLine(ESCAPE_AS_ARG_PAT, line);

                Node node = getNode(m.group("nodeIdx"), idxToNode, line);
                String callNode = m.group("escapeNode");

                propagations.put(node, new Reason.EscapesAsArg(callNode));
            } else if (line.contains("merged with other object")) {
                Matcher m = parseLine(MERGED_PAT, line);

                Node node = getNode(m.group("nodeIdx"), idxToNode, line);
                Node srcNode = getNode(m.group("srcNodeIdx"), idxToNode, line);

                propagations.put(node, new Reason.MergedWithObject(srcNode));
            } else if (line.contains("has field with unknown offset")) {
                Matcher m = parseLine(UNKNOWN_FIELD_OFFSET_PAT, line);

                Node node = getNode(m.group("nodeIdx"), idxToNode, line);

                propagations.put(node, new Reason.HasFieldWithUnknownOffset());
            } else if (line.contains("blackhole")) {
                Matcher m = parseLine(BLACKHOLE_PAT, line);

                Node node = getNode(m.group("nodeIdx"), idxToNode, line);

                propagations.put(node, new Reason.BlackHole());
            } else {
                throw new ParseError("Unrecognized line: " + line);
            }
        }
        return new EAInvocation(List.copyOf(allocations), Map.copyOf(propagations));
    }

    private static Matcher parseLine(Pattern pat, String line) {
        Matcher m = pat.matcher(line);
        if (!m.find()) {
            throw new ParseError("Can not parse line: " + line);
        }
        return m;
    }

    private static Node getNode(String idxGroup, Map<Integer, Node> idxToNode, String line) {
        int nodeIdx = Integer.parseInt(idxGroup);
        Node node = Objects.requireNonNull(idxToNode.get(nodeIdx), "Could not find node: " + nodeIdx + ", for line: " + line);
        return node;
    }

    enum EscapeType { NoEscape, ArgEscape, GlobalEscape }
    enum NodeType { JavaObject, LocalVar, Field }

    record Node(NodeType type, int idx, String fullDump) {
        @Override
        public String toString() {
            return  header() + " " + fullDump;
        }

        public String header() {
            return type() + "(" + idx() + ")";
        }
    }

    static sealed interface Reason {
        record PropagatedFromNode(Node node) implements Reason {}
        record EscapesAsArg(String callNode) implements Reason {
            @Override
            public String simpleReason() {
                Matcher matcher = FIRST_CALL_PAT.matcher(callNode());
                matcher.find();

                return "Escapes as argument to call to: " + matcher.group("callname");
            }
        }
        record MergedWithObject(Node other) implements Reason {}
        record HasFieldWithUnknownOffset() implements Reason {}
        record BlackHole() implements Reason {}

        default String simpleReason() { return toString(); }
    }

    record EAInvocation(List<Node> allocations, Map<Node, Reason> propagations) {}

    static class ParseError extends Error {
        public ParseError(String message) {
            super(message);
        }
    }
}