/*

 * 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.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * 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.

 */

package indify;

import java.util.*;

import java.io.*;

import java.lang.reflect.Modifier;

import java.util.regex.*;

/**

 * Transform one or more class files to incorporate JSR 292 features,

 * such as {@code invokedynamic}.

 * <p>

 * This is a standalone program in a single source file.

 * In this form, it may be useful for test harnesses, small experiments, and javadoc examples.

 * Copies of this file may show up in multiple locations for standalone usage.

 * The primary maintained location of this file is as follows:

 * <a href="http://kenai.com/projects/ninja/sources/indify-repo/content/src/indify/Indify.java">

 * http://kenai.com/projects/ninja/sources/indify-repo/content/src/indify/Indify.java</a>

 * <p>

 * Static private methods named MH_x and MT_x (where x is arbitrary)

 * must be stereotyped generators of MethodHandle and MethodType

 * constants.  All calls to them are transformed to {@code CONSTANT_MethodHandle}

 * and {@code CONSTANT_MethodType} "ldc" instructions.

 * The stereotyped code must create method types by calls to {@code methodType} or

 * {@code fromMethodDescriptorString}.  The "lookup" argument must be created

 * by calls to {@code java.lang.invoke.MethodHandles#lookup MethodHandles.lookup}.

 * The class and string arguments must be constant.

 * The following methods of {@code java.lang.invoke.MethodHandle.Lookup Lookup} are

 * allowed for method handle creation: {@code findStatic}, {@code findVirtual},

 * {@code findConstructor}, {@code findSpecial},

 * {@code findGetter}, {@code findSetter},

 * {@code findStaticGetter}, or {@code findStaticSetter}.

 * The call to one of these methods must be followed immediately

 * by an {@code areturn} instruction.

 * The net result of the call to the MH_x or MT_x method must be

 * the creation of a constant method handle.  Thus, replacing calls

 * to MH_x or MT_x methods by {@code ldc} instructions should leave

 * the meaning of the program unchanged.

 * <p>

 * Static private methods named INDY_x must be stereotyped generators

 * of {@code invokedynamic} call sites.

 * All calls to them must be immediately followed by

 * {@code invokeExact} calls.

 * All such pairs of calls are transformed to {@code invokedynamic}

 * instructions.  Each INDY_x method must begin with a call to a

 * MH_x method, which is taken to be its bootstrap method.

 * The method must be immediately invoked (via {@code invokeGeneric}

 * on constant lookup, name, and type arguments.  An object array of

 * constants may also be appended to the {@code invokeGeneric call}.

 * This call must be cast to {@code CallSite}, and the result must be

 * immediately followed by a call to {@code dynamicInvoker}, with the

 * resulting method handle returned.

 * <p>

 * The net result of all of these actions is equivalent to the JVM's

 * execution of an {@code invokedynamic} instruction in the unlinked state.

 * Running this code once should produce the same results as running

 * the corresponding {@code invokedynamic} instruction.

 * In order to model the caching behavior, the code of an INDY_x

 * method is allowed to begin with getstatic, aaload, and if_acmpne

 * instructions which load a static method handle value and return it

 * if the value is non-null.

 * <p>

 * Example usage:

 * <blockquote><pre>

$ JAVA_HOME=(some recent OpenJDK 7 build)

$ ant

$ $JAVA_HOME/bin/java -cp build/classes indify.Indify --overwrite --dest build/testout build/classes/indify/Example.class

$ $JAVA_HOME/bin/java -XX:+UnlockExperimentalVMOptions -XX:+EnableInvokeDynamic -cp build/classes indify.Example

MT = (java.lang.Object)java.lang.Object

MH = adder(int,int)java.lang.Integer

adder(1,2) = 3

calling indy:  42

$ $JAVA_HOME/bin/java -XX:+UnlockExperimentalVMOptions -XX:+EnableInvokeDynamic -cp build/testout indify.Example

(same output as above)

 * </pre></blockquote>

 * <p>

 * Before OpenJDK build b123, the format of {@code CONSTANT_InvokeDynamic} is in transition,

 * and the switch {@code --transitionalJSR292=yes} is recommended.

 * It is turned <em>off</em> by default, but users of earlier builds may need to turn it on.

 * <p>

 * A version of this transformation built on top of <a href="http://asm.ow2.org/">http://asm.ow2.org/</a> would be welcome.

 * @author John Rose

 */

public class Indify {

    public static void main(String... av) throws IOException {

        new Indify().run(av);

    }

    public File dest;

    public String[] classpath = {"."};

    public boolean keepgoing = false;

    public boolean expandProperties = false;

    public boolean overwrite = false;

    public boolean quiet = false;

    public boolean verbose = false;

    public boolean transitionalJSR292 = false;  // final version is distributed

    public boolean all = false;

    public int verifySpecifierCount = -1;

    public void run(String... av) throws IOException {

        List<String> avl = new ArrayList<>(Arrays.asList(av));

        parseOptions(avl);

        if (avl.isEmpty())

            throw new IllegalArgumentException("Usage: indify [--dest dir] [option...] file...");

        if ("--java".equals(avl.get(0))) {

            avl.remove(0);

            try {

                runApplication(avl.toArray(new String[0]));

            } catch (Exception ex) {

                if (ex instanceof RuntimeException)  throw (RuntimeException) ex;

                throw new RuntimeException(ex);

            }

            return;

        }

        Exception err = null;

        for (String a : avl) {

            try {

                indify(a);

            } catch (Exception ex) {

                if (err == null)  err = ex;

                System.err.println("failure on "+a);

                if (!keepgoing)  break;

            }

        }

        if (err != null) {

            if (err instanceof IOException)  throw (IOException) err;

            throw (RuntimeException) err;

        }

    }

    /** Execute the given application under a class loader which indifies all application classes. */

    public void runApplication(String... av) throws Exception {

        List<String> avl = new ArrayList<>(Arrays.asList(av));

        String mainClassName = avl.remove(0);

        av = avl.toArray(new String[0]);

        Class<?> mainClass = Class.forName(mainClassName, true, makeClassLoader());

        java.lang.reflect.Method main = mainClass.getMethod("main", String[].class);

        try { main.setAccessible(true); } catch (SecurityException ex) { }

        main.invoke(null, (Object) av);

    }

    public void parseOptions(List<String> av) throws IOException {

        for (; !av.isEmpty(); av.remove(0)) {

            String a = av.get(0);

            if (a.startsWith("-")) {

                String a2 = null;

                int eq = a.indexOf('=');

                if (eq > 0) {

                    a2 = maybeExpandProperties(a.substring(eq+1));

                    a = a.substring(0, eq+1);

                }

                switch (a) {

                case "--java":

                    return;  // keep this argument

                case "-d": case "--dest": case "-d=": case "--dest=":

                    dest = new File(a2 != null ? a2 : maybeExpandProperties(av.remove(1)));

                    break;

                case "-cp": case "--classpath":

                    classpath = maybeExpandProperties(av.remove(1)).split("["+File.pathSeparatorChar+"]");

                    break;

                case "-k": case "--keepgoing": case "--keepgoing=":

                    keepgoing = booleanOption(a2);  // print errors but keep going

                    break;

                case "--expand-properties": case "--expand-properties=":

                    expandProperties = booleanOption(a2);  // expand property references in subsequent arguments

                    break;

                case "--verify-specifier-count": case "--verify-specifier-count=":

                    verifySpecifierCount = Integer.valueOf(a2);

                    break;

                case "--overwrite": case "--overwrite=":

                    overwrite = booleanOption(a2);  // overwrite output files

                    break;

                case "--all": case "--all=":

                    all = booleanOption(a2);  // copy all classes, even if no patterns

                    break;

                case "-q": case "--quiet": case "--quiet=":

                    quiet = booleanOption(a2);  // less output

                    break;

                case "-v": case "--verbose": case "--verbose=":

                    verbose = booleanOption(a2);  // more output

                    break;

                case "--transitionalJSR292": case "--transitionalJSR292=":

                    transitionalJSR292 = booleanOption(a2);  // use older invokedynamic format

                    break;

                default:

                    throw new IllegalArgumentException("unrecognized flag: "+a);

                }

                continue;

            } else {

                break;

            }

        }

        if (dest == null && !overwrite)

            throw new RuntimeException("no output specified; need --dest d or --overwrite");

        if (expandProperties) {

            for (int i = 0; i < av.size(); i++)

                av.set(i, maybeExpandProperties(av.get(i)));

        }

    }

    private boolean booleanOption(String s) {

        if (s == null)  return true;

        switch (s) {

        case "true":  case "yes": case "on":  case "1": return true;

        case "false": case "no":  case "off": case "0": return false;

        }

        throw new IllegalArgumentException("unrecognized boolean flag="+s);

    }

    private String maybeExpandProperties(String s) {

        if (!expandProperties)  return s;

        Set<String> propsDone = new HashSet<>();

        while (s.contains("${")) {

            int lbrk = s.indexOf("${");

            int rbrk = s.indexOf('}', lbrk);

            if (rbrk < 0)  break;

            String prop = s.substring(lbrk+2, rbrk);

            if (!propsDone.add(prop))  break;

            String value = System.getProperty(prop);

            if (verbose)  System.err.println("expanding ${"+prop+"} => "+value);

            if (value == null)  break;

            s = s.substring(0, lbrk) + value + s.substring(rbrk+1);

        }

        return s;

    }

    public void indify(String a) throws IOException {

        File f = new File(a);

        String fn = f.getName();

        if (fn.endsWith(".class") && f.isFile())

            indifyFile(f, dest);

        else if (fn.endsWith(".jar") && f.isFile())

            indifyJar(f, dest);

        else if (f.isDirectory())

            indifyTree(f, dest);

        else if (!keepgoing)

            throw new RuntimeException("unrecognized file: "+a);

    }

    private void ensureDirectory(File dir) {

        if (dir.mkdirs() && !quiet)

            System.err.println("created "+dir);

    }

    public void indifyFile(File f, File dest) throws IOException {

        if (verbose)  System.err.println("reading "+f);

        ClassFile cf = new ClassFile(f);

        Logic logic = new Logic(cf);

        boolean changed = logic.transform();

        logic.reportPatternMethods(quiet, keepgoing);

        if (changed || all) {

            File outfile;

            if (dest != null) {

                ensureDirectory(dest);

                outfile = classPathFile(dest, cf.nameString());

            } else {

                outfile = f;  // overwrite input file, no matter where it is

            }

            cf.writeTo(outfile);

            if (!quiet)  System.err.println("wrote "+outfile);

        }

    }

    File classPathFile(File pathDir, String className) {

        String qualname = className.replace('.','/')+".class";

        qualname = qualname.replace('/', File.separatorChar);

        return new File(pathDir, qualname);

    }

    public void indifyJar(File f, Object dest) throws IOException {

        throw new UnsupportedOperationException("Not yet implemented");

    }

    public void indifyTree(File f, File dest) throws IOException {

        if (verbose)  System.err.println("reading directory: "+f);

        for (File f2 : f.listFiles(new FilenameFilter() {

                public boolean accept(File dir, String name) {

                    if (name.endsWith(".class"))  return true;

                    if (name.contains("."))  return false;

                    // return true if it might be a package name:

                    return Character.isJavaIdentifierStart(name.charAt(0));

                }})) {

            if (f2.getName().endsWith(".class"))

                indifyFile(f2, dest);

            else if (f2.isDirectory())

                indifyTree(f2, dest);

        }

    }

    public ClassLoader makeClassLoader() {

        return new Loader();

    }

    private class Loader extends ClassLoader {

        Loader() {

            this(Indify.class.getClassLoader());

        }

        Loader(ClassLoader parent) {

            super(parent);

        }

        public Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {

            File f = findClassInPath(name);

            if (f != null) {

                try {

                    Class<?> c = transformAndLoadClass(f);

                    if (c != null) {

                        if (resolve)  resolveClass(c);

                        return c;

                    }

                } catch (Exception ex) {

                    if (ex instanceof IllegalArgumentException)

                        // pass error from reportPatternMethods

                        throw (IllegalArgumentException) ex;

                }

            }

            return super.loadClass(name, resolve);

        }

        private File findClassInPath(String name) {

            for (String s : classpath) {

                File f = classPathFile(new File(s), name);

                //System.out.println("Checking for "+f);

                if (f.exists() && f.canRead()) {

                    return f;

                }

            }

            return null;

        }

        protected Class<?> findClass(String name) throws ClassNotFoundException {

            try {

                File f = findClassInPath(name);

                if (f != null) {

                    Class<?> c = transformAndLoadClass(f);

                    if (c != null)  return c;

                }

            } catch (IOException ex) {

                throw new ClassNotFoundException("IO error", ex);

            }

            throw new ClassNotFoundException();

        }

        private Class<?> transformAndLoadClass(File f) throws ClassNotFoundException, IOException {

            if (verbose)  System.err.println("Loading class from "+f);

            ClassFile cf = new ClassFile(f);

            Logic logic = new Logic(cf);

            boolean changed = logic.transform();

            if (verbose && !changed)  System.err.println("(no change)");

            logic.reportPatternMethods(!verbose, keepgoing);

            byte[] bytes = cf.toByteArray();

            return defineClass(null, bytes, 0, bytes.length);

        }

    }

    private class Logic {

        // Indify logic, per se.

        ClassFile cf;

        final char[] poolMarks;

        final Map<Method,Constant> constants = new HashMap<>();

        final Map<Method,String> indySignatures = new HashMap<>();

        Logic(ClassFile cf) {

            this.cf = cf;

            poolMarks = new char[cf.pool.size()];

        }

        boolean transform() {

            if (!initializeMarks())  return false;

            if (!findPatternMethods())  return false;

            Pool pool = cf.pool;

            //for (Constant c : cp)  System.out.println("  # "+c);

            for (Method m : cf.methods) {

                if (constants.containsKey(m))  continue;  // don't bother

                // Transform references.

                int blab = 0;

                for (Instruction i = m.instructions(); i != null; i = i.next()) {

                    if (i.bc != opc_invokestatic)  continue;

                    int methi = i.u2At(1);

                    if (poolMarks[methi] == 0)  continue;

                    Short[] ref = pool.getMemberRef((short)methi);

                    Method conm = findMember(cf.methods, ref[1], ref[2]);

                    if (conm == null)  continue;

                    Constant con = constants.get(conm);

                    if (con == null)  continue;

                    if (blab++ == 0 && !quiet)

                        System.err.println("patching "+cf.nameString()+"."+m);

                    //if (blab == 1) { for (Instruction j = m.instructions(); j != null; j = j.next()) System.out.println("  |"+j); }

                    if (con.tag == CONSTANT_InvokeDynamic ||

                        con.tag == CONSTANT_InvokeDynamic_17) {

                        // need to patch the following instruction too,

                        // but there are usually intervening argument pushes too

                        Instruction i2 = findPop(i);

                        Short[] ref2 = null;

                        short ref2i = 0;

                        if (i2 != null && i2.bc == opc_invokevirtual &&

                                poolMarks[(char)(ref2i = (short) i2.u2At(1))] == 'D')

                            ref2 = pool.getMemberRef(ref2i);

                        if (ref2 == null || !"invokeExact".equals(pool.getString(ref2[1]))) {

                            System.err.println(m+": failed to create invokedynamic at "+i.pc);

                            continue;

                        }

                        String invType = pool.getString(ref2[2]);

                        String bsmType = indySignatures.get(conm);

                        if (!invType.equals(bsmType)) {

                            System.err.println(m+": warning: "+conm+" call type and local invoke type differ: "

                                    +bsmType+", "+invType);

                        }

                        assert(i.len == 3 || i2.len == 3);

                        if (!quiet)  System.err.println(i+" "+conm+";...; "+i2+" => invokedynamic "+con);

                        int start = i.pc + 3, end = i2.pc;

                        System.arraycopy(i.codeBase, start, i.codeBase, i.pc, end-start);

                        i.forceNext(0);  // force revisit of new instruction

                        i2.u1AtPut(-3, opc_invokedynamic);

                        i2.u2AtPut(-2, con.index);

                        i2.u2AtPut(0, (short)0);

                        i2.u1AtPut(2, opc_nop);

                        //System.out.println(new Instruction(i.codeBase, i2.pc-3));

                    } else {

                        if (!quiet)  System.err.println(i+" "+conm+" => ldc "+con);

                        assert(i.len == 3);

                        i.u1AtPut(0, opc_ldc_w);

                        i.u2AtPut(1, con.index);

                    }

                }

                //if (blab >= 1) { for (Instruction j = m.instructions(); j != null; j = j.next()) System.out.println("    |"+j); }

            }

            cf.methods.removeAll(constants.keySet());

            return true;

        }

        // Scan forward from the instruction to find where the stack p

        // below the current sp at the instruction.

        Instruction findPop(Instruction i) {

            //System.out.println("findPop from "+i);

            Pool pool = cf.pool;

            JVMState jvm = new JVMState();

        decode:

            for (i = i.clone().next(); i != null; i = i.next()) {

                String pops = INSTRUCTION_POPS[i.bc];

                //System.out.println("  "+i+" "+jvm.stack+" : "+pops.replace("$", " => "));

                if (pops == null)  break;

                if (jvm.stackMotion(i.bc))  continue decode;

                if (pops.indexOf('Q') >= 0) {

                    Short[] ref = pool.getMemberRef((short) i.u2At(1));

                    String type = simplifyType(pool.getString(CONSTANT_Utf8, ref[2]));

                    switch (i.bc) {

                    case opc_getstatic:

                    case opc_getfield:

                    case opc_putstatic:

                    case opc_putfield:

                        pops = pops.replace("Q", type);

                        break;

                    default:

                        if (!type.startsWith("("))

                            throw new InternalError(i.toString());

                        pops = pops.replace("Q$Q", type.substring(1).replace(")","$"));

                        break;

                    }

                    //System.out.println("special type: "+type+" => "+pops);

                }

                int npops = pops.indexOf('$');

                if (npops < 0)  throw new InternalError();

                if (npops > jvm.sp())  return i;

                List<Object> args = jvm.args(npops);

                int k = 0;

                for (Object x : args) {

                    char have = (Character) x;

                    char want = pops.charAt(k++);

                    if (have == 'X' || want == 'X')  continue;

                    if (have != want)  break decode;

                }

                if (pops.charAt(k++) != '$')  break decode;

                args.clear();

                while (k < pops.length())

                    args.add(pops.charAt(k++));

            }

            System.err.println("*** bailout on jvm: "+jvm.stack+" "+i);

            return null;

        }

        boolean findPatternMethods() {

            boolean found = false;

            for (char mark : "THI".toCharArray()) {

                for (Method m : cf.methods) {

                    if (!Modifier.isPrivate(m.access))  continue;

                    if (!Modifier.isStatic(m.access))  continue;

                    if (nameAndTypeMark(m.name, m.type) == mark) {

                        Constant con = scanPattern(m, mark);

                        if (con == null)  continue;

                        constants.put(m, con);

                        found = true;

                    }

                }

            }

            return found;

        }

        void reportPatternMethods(boolean quietly, boolean allowMatchFailure) {

            if (!quietly && !constants.keySet().isEmpty())

                System.err.println("pattern methods removed: "+constants.keySet());

            for (Method m : cf.methods) {

                if (nameMark(cf.pool.getString(m.name)) != 0 &&

                    constants.get(m) == null) {

                    String failure = "method has special name but fails to match pattern: "+m;

                    if (!allowMatchFailure)

                        throw new IllegalArgumentException(failure);

                    else if (!quietly)