IO.java

package org.rascalmpl.library.lang.xml;

import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
import java.io.Writer;
import java.net.MalformedURLException;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import org.jsoup.Jsoup;
import org.jsoup.internal.SimpleStreamReader;
import org.jsoup.nodes.Attribute;
import org.jsoup.nodes.Attributes;
import org.jsoup.nodes.Comment;
import org.jsoup.nodes.DataNode;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Document.OutputSettings;
import org.jsoup.nodes.Document.OutputSettings.Syntax;
import org.jsoup.nodes.DocumentType;
import org.jsoup.nodes.Element;
import org.jsoup.nodes.Entities.EscapeMode;
import org.jsoup.nodes.Node;
import org.jsoup.nodes.Range;
import org.jsoup.nodes.TextNode;
import org.jsoup.nodes.XmlDeclaration;
import org.jsoup.parser.ParseSettings;
import org.jsoup.parser.Parser;
import org.jsoup.parser.StreamParser;
import org.rascalmpl.exceptions.RuntimeExceptionFactory;
import org.rascalmpl.uri.URIResolverRegistry;
import org.rascalmpl.values.IRascalValueFactory;
import org.rascalmpl.values.functions.IFunction;
import org.rascalmpl.values.maybe.UtilMaybe;

import io.usethesource.vallang.IBool;
import io.usethesource.vallang.IConstructor;
import io.usethesource.vallang.IDateTime;
import io.usethesource.vallang.IExternalValue;
import io.usethesource.vallang.IInteger;
import io.usethesource.vallang.IList;
import io.usethesource.vallang.IMap;
import io.usethesource.vallang.IMapWriter;
import io.usethesource.vallang.INode;
import io.usethesource.vallang.IRational;
import io.usethesource.vallang.IReal;
import io.usethesource.vallang.ISet;
import io.usethesource.vallang.ISourceLocation;
import io.usethesource.vallang.IString;
import io.usethesource.vallang.ITuple;
import io.usethesource.vallang.IValue;
import io.usethesource.vallang.type.Type;
import io.usethesource.vallang.type.TypeFactory;
import io.usethesource.vallang.visitors.IValueVisitor;

public class IO {
    private final IRascalValueFactory vf;
    private static final String SRC_ATTR = "src";
    private static final String QUALIFIED_SRC_ATTR = "rascal-src";
    private static final String SRCS_ATTR = "srcs";
    private static final String QUALIFIED_SRCS_ATTR = "rascal-srcs";

    private static final TypeFactory tf = TypeFactory.getInstance();
    // `Maybe[value] ()`
    private static final Type nextFunctionType = tf.functionType(UtilMaybe.Maybe.instantiate(Map.of(UtilMaybe.ParameterT, tf.valueType())), tf.tupleEmpty(), tf.tupleEmpty());

    public IO(IRascalValueFactory vf) {
        this.vf = vf;
    }
    
    public IValue readXML(ISourceLocation loc, IBool fullyQualify, IBool trackOrigins, IBool includeEndTags,  IBool ignoreComments, IBool ignoreWhitespace, IString charset, IBool inferCharset) {
        if (inferCharset.getValue()) {
            charset = vf.string(URIResolverRegistry.getInstance().detectCharset(loc).toString());
        }
       
        try (InputStream reader = URIResolverRegistry.getInstance().getInputStream(loc)) {
            Parser xmlParser = Parser.xmlParser()
                .settings(new ParseSettings(true, true))
                .setTrackPosition(trackOrigins.getValue())
                ;
            
            Document doc = Jsoup.parse(reader, charset.getValue(), loc.getURI().toString(), xmlParser);
            
            return toINode(doc, trackOrigins.getValue() ? loc : null, fullyQualify.getValue(), includeEndTags.getValue(), ignoreWhitespace.getValue(), ignoreComments.getValue());
        } 
        catch (MalformedURLException e) {
            throw RuntimeExceptionFactory.malformedURI(loc.getURI().toASCIIString());
        } 
        catch (IOException e) {
            throw RuntimeExceptionFactory.io(e);
        }
    }

    public IFunction streamXML(ISourceLocation loc, IString elementName, IBool fullyQualify, IBool trackOrigins, IBool includeEndTags, IBool ignoreComments, IBool ignoreWhitespace, IString charset, IBool inferCharset) {
        if (inferCharset.getValue()) {
            charset = vf.string(URIResolverRegistry.getInstance().detectCharset(loc).toString());
        }

        try {
            InputStream reader = URIResolverRegistry.getInstance().getInputStream(loc);
            Parser xmlParser = Parser.xmlParser()
                .settings(new ParseSettings(false, false))
                .setTrackPosition(trackOrigins.getValue());   

            StreamParser streamer = new StreamParser(xmlParser);
            Reader ssr = new SimpleStreamReader(reader, Charset.forName(charset.getValue()));
            
            streamer.parse(ssr, loc.getURI().toString());

            return vf.function(nextFunctionType, (args, kwargs) -> {
                try {
                    var elem = streamer.selectNext(elementName.getValue());

                    if (elem == null) {
                        // we're done
                        try {
                            reader.close();
                            streamer.close();
                        } catch (IOException e) { /* ignored */}

                        return UtilMaybe.nothing();
                    }
                    else {
                        try {
                            return UtilMaybe.just(toINode(elem, trackOrigins.getValue() ? loc : null, fullyQualify.getValue(), includeEndTags.getValue(), ignoreWhitespace.getValue(), ignoreComments.getValue()));
                        }
                        finally {
                            // clean up to safe memory space just before the return jump
                            elem.remove();
                        }     
                    }
                }
                catch (IOException e) {
                    throw RuntimeExceptionFactory.io(e);
                }
            });        
        }
        catch (IOException e) {
            throw RuntimeExceptionFactory.io(e);
        }
    }


    public IValue readXML(IString string, ISourceLocation src, IBool fullyQualify, IBool trackOrigins, IBool includeEndTags, IBool ignoreComments, IBool ignoreWhitespace) {
        if (string.length() == 0) {
            throw RuntimeExceptionFactory.io("empty XML document");
        }

        Parser xmlParser = Parser.xmlParser()
                .settings(new ParseSettings(false, false))
                .setTrackPosition(trackOrigins.getValue())
                ;
             
        Document doc = Jsoup.parse(string.getValue(), src.getURI().toString(), xmlParser);
            
        return toINode(doc.firstChild(), trackOrigins.getValue() ? src : null, fullyQualify.getValue(),  includeEndTags.getValue(), ignoreWhitespace.getValue(), ignoreComments.getValue());        
    }

    private IValue toINode(Document doc, ISourceLocation file, boolean fullyQualify, boolean includeEndTags, boolean ignoreWhitespace, boolean ignoreComments) {
        return toINode((Node) doc, file, fullyQualify, includeEndTags, ignoreWhitespace, ignoreComments);
    }

    private IValue toINode(Node node, ISourceLocation file, boolean fullyQualify, boolean includeEndTags, boolean ignoreWhitespace, boolean ignoreComments) {
        if (node instanceof TextNode) {
            return toIString((TextNode) node, file);
        }
        else if (node instanceof DocumentType) {
            DocumentType dt = (DocumentType) node;
            Map<String, IValue> args = new HashMap<>();
            args.put("publicId", vf.string(dt.publicId()));
            args.put("systemId", vf.string(dt.systemId()));
            return vf.node("documentType", new IValue[0], args);
        }
        else if (node instanceof XmlDeclaration) {
            XmlDeclaration xd = (XmlDeclaration) node;
            Map<String, IValue> args = new HashMap<>();
            args.put("content", vf.string(xd.getWholeDeclaration()));
            return vf.node("xmlDeclaration", new IValue[0], args);
        }
        else if (node instanceof DataNode) {
            return toIString((DataNode) node, file);
        }
        else if (node instanceof Comment) {
            return vf.node("comment", vf.string(((Comment) node).getData()));
        }
        else if (node instanceof Element) {            
            Element elem = (Element) node;

            IMapWriter namespaces = vf.mapWriter();
            
            Map<String,IValue> kws = 
                StreamSupport.stream(elem.attributes().spliterator(), false)
                    .filter(a -> { 
                        // collect the namespaces in a map
                        if (a.getKey().startsWith("xmlns:")) {
                            namespaces.put(vf.string(a.getKey().substring("xmlns:".length())), vf.string(a.getValue()));
                        } 
                        else if (a.getKey().equals("xmlns")) {
                            namespaces.put(vf.string("xml"), vf.string(a.getValue()));
                        }

                        // remove all the namespace attributes
                        return !a.getKey().startsWith("xmlns");
                    })
                    .collect(Collectors.toMap(a -> normalizeAttr(a.getKey()), a -> vf.string(removeNamespace(a, elem.attributes(), fullyQualify))));
            
            // we traverse again to record the source positions of each attribute
            IMap Srcs = file != null ? StreamSupport.stream(elem.attributes().spliterator(), false)
                    .filter(a -> !a.getKey().startsWith("xmlns"))
                    .map(a -> vf.tuple(vf.string(normalizeAttr(removeNamespace(a, elem.attributes(), fullyQualify))), attrToLoc(a, file))) // key-value tuples for the map collector
                    .filter(t -> !t.get(1).equals(vf.tuple())) // remove the failed location lookups for robustness sake
                    .collect(vf.mapWriter())
                : vf.map()
            ;

            if (fullyQualify) {
                IMap m = namespaces.done();

                if (m.size() > 0) {
                    kws.put("xmlns", m);
                }
            }

            IValue[] args = elem.childNodes().stream()
                .filter(p -> !ignoreComments || (!(p instanceof Comment) && !(p instanceof DocumentType)))
                .filter(p -> !ignoreWhitespace || !(p instanceof TextNode && ((TextNode) p).isBlank()))
                .map(n -> toINode(n, file, fullyQualify, includeEndTags, ignoreWhitespace, ignoreComments))
                .filter(c -> c != null) // filter document types
                .toArray(IValue[]::new);

            if (file != null) {
                kws.put(kws.containsKey(SRC_ATTR) ? QUALIFIED_SRC_ATTR : SRC_ATTR, nodeToLoc((Element) node, file, includeEndTags));
                if (!Srcs.isEmpty()) {
                    kws.put(kws.containsKey(SRCS_ATTR) ? QUALIFIED_SRCS_ATTR : SRCS_ATTR, Srcs);
                }
            }

            return vf.node(removeNamespace(node.nodeName(), fullyQualify), args).asWithKeywordParameters().setParameters(kws);
        }
        else {
            throw RuntimeExceptionFactory.illegalArgument(vf.string(node.toString()), vf.string("unexpected kind of XML node: " + node.getClass().getCanonicalName()));
        }
    }

    private static String removeNamespace(String name, boolean fullyQualify) {
        if (fullyQualify) {
            return name;
        }

        int index = name.indexOf(":");

        if (index == -1) {
            return name;
        }

        return name.substring(index+1);

    }
    private static String removeNamespace(Attribute a, Attributes otherAttributes, boolean fullyQualify) {
        if (fullyQualify) {
            return a.getKey();
        }
        
        String key = a.getKey();
        int index = key.indexOf(":");

        if (index == -1) {
            return a.getKey();
        }

        return key.substring(index+1);
    }

    private ITuple attrToLoc(Attribute a, ISourceLocation file) {
        Range nameRange = a.sourceRange().nameRange();
        Range valueRange = a.sourceRange().valueRange();

        if (valueRange.start().pos() < 0 || nameRange.start().pos() < 0) {
            // this is strange, tagging an error here so it can be filtered. 
            assert false;
            return vf.tuple();
        }

        return vf.tuple(
            vf.sourceLocation(file, 
            nameRange.start().pos(), 
            nameRange.end().pos() - nameRange.start().pos(),
            nameRange.start().lineNumber(),
            nameRange.end().lineNumber(),
            nameRange.start().columnNumber() - 1,
            nameRange.end().columnNumber() - 1)
            , vf.sourceLocation(file, 
            valueRange.start().pos(), 
            valueRange.end().pos() - valueRange.start().pos(),
            valueRange.start().lineNumber(),
            valueRange.end().lineNumber(),
            valueRange.start().columnNumber() - 1,
            valueRange.end().columnNumber() - 1)
        );
    }

    private ISourceLocation nodeToLoc(Element node, ISourceLocation file, boolean includeEndTags) {
        Range startRange = node.sourceRange();
        if (!startRange.isTracked()) {
            return file;
        }
        
        Range endRange = node.endSourceRange();

        return includeEndTags && endRange.isTracked()
            ? vf.sourceLocation(file,
                startRange.start().pos(),
                endRange.end().pos() - startRange.start().pos(),
                startRange.start().lineNumber(),
                endRange.end().lineNumber(),
                startRange.start().columnNumber() - 1,
                endRange.end().columnNumber() - 1
            )
            : vf.sourceLocation(file, 
                startRange.start().pos(), 
                startRange.end().pos() - startRange.start().pos(),
                startRange.start().lineNumber(),
                startRange.end().lineNumber(),
                startRange.start().columnNumber() - 1,
                startRange.end().columnNumber() - 1
                );
    }

    private IValue toIString(DataNode node, ISourceLocation file) {
        return vf.string(node.getWholeData());
    }

    private IValue toIString(TextNode elem, ISourceLocation file) {
        return vf.string(elem.getWholeText());
    }

    /**
     * Produces a XML string output from an arbitrary Rascal value
     * 
     * Why go through all the trouble of building a DOM? The only reason is compliance.
     * Escapes, encodings, etc. all are maintained by these classes from jsoup
     */
    public IString writeXMLString(IValue cons, IString charset, IBool outline, IBool prettyPrint, IInteger indentAmount, IInteger maxPaddingWidth, IBool dropOrigins) {
        try {
            Document doc = createXMLDocument(cons, dropOrigins.getValue());
            doc = doc.outputSettings(createOutputSettings(charset.getValue(), outline.getValue(), prettyPrint.getValue(), indentAmount.intValue(), maxPaddingWidth.intValue()));
            
            return vf.string(doc.outerHtml());
        }
        catch (IOException e) {
            throw RuntimeExceptionFactory.io(e);
        }
    }

    /**
     * Produces a XML string output from an node value or algebraic data-type value.
     * 
     * Why go through all the trouble of building a DOM? The only reason is compliance.
     * Escapes, encodings, etc. all are maintained by these classes from JSoup
     */
    public void writeXMLFile(ISourceLocation file, IValue cons, IString charset, IBool outline, IBool prettyPrint, IInteger indentAmount, IInteger maxPaddingWidth, IBool dropOrigins) {
        
        try (Writer out = URIResolverRegistry.getInstance().getCharacterWriter(file, charset.getValue(), false)) {
            Document doc = createXMLDocument(cons, dropOrigins.getValue());
            doc = doc.outputSettings(createOutputSettings(charset.getValue(), outline.getValue(), prettyPrint.getValue(), indentAmount.intValue(), maxPaddingWidth.intValue()));
            out.write(doc.outerHtml());
        }
        catch (IOException e) {
            throw RuntimeExceptionFactory.io(e);
        }
    }

    private OutputSettings createOutputSettings(String charset, boolean outline, boolean prettyPrint, int indentAmount, int maxPaddingWidth) {
        return new OutputSettings()
            .charset(charset)
            .escapeMode(EscapeMode.base)
            .outline(outline)
            .prettyPrint(prettyPrint)
            .indentAmount(indentAmount)
            .maxPaddingWidth(maxPaddingWidth)
            .syntax(Syntax.xml);
    }

    /**
     * Translates a constructor tree to a xml DOM tree
     */
    private Document createXMLDocument(IValue cons, boolean dropOrigins) throws IOException {
        Document doc = new Document("http://localhost");

        Node node = cons.accept(new ElementCreator(dropOrigins));
        
        return (Document) doc.appendChild(node);
    }

    private static class ElementCreator implements IValueVisitor<Node, RuntimeException> {
        private final boolean dropOrigins;

        public ElementCreator(boolean dropOrigins) {
            this.dropOrigins = dropOrigins;
        }

        @Override
        public Node visitString(IString o) throws RuntimeException {
            return new TextNode(o.getValue());
        }

        @Override
        public Node visitReal(IReal o) throws RuntimeException {
            Element real = new Element("real");
            real.attr("val", o.toString());
            return real;
        }

        @Override
        public Node visitRational(IRational o) throws RuntimeException {
            Element rat = new Element("rat");
            rat.attr("numerator", o.numerator().toString());
            rat.attr("denominator", o.denominator().toString());
            return rat;
        }

        @Override
        public Node visitList(IList o) throws RuntimeException {
            Element list = new Element("list");
            o.stream().forEach(e -> list.appendChild(e.accept(this)));
            return list;
        }

        @Override
        public Node visitSet(ISet o) throws RuntimeException {
            Element list = new Element("set");
            o.stream().forEach(e -> list.appendChild(e.accept(this)));
            return list;
        }

        @Override
        public Node visitSourceLocation(ISourceLocation o) throws RuntimeException {
            Element loc = new Element("loc");

            loc.attr("uri", o.getURI().toString());

            if (o.hasOffsetLength()) {
                loc.attr("offset", Integer.toString(o.getOffset()));
                loc.attr("length", Integer.toString(o.getLength()));
            }

            if (o.hasLineColumn()) {
                loc.attr("beginLine", Integer.toString(o.getBeginLine()));
                loc.attr("endLine", Integer.toString(o.getEndLine()));
                loc.attr("beginColumn", Integer.toString(o.getBeginColumn()));
                loc.attr("endColumn", Integer.toString(o.getEndColumn()));
            }

            return loc;
        }

        @Override
        public Node visitTuple(ITuple o) throws RuntimeException {
            Element tuple = new Element("tuple");
            StreamSupport.stream(o.spliterator(), false)
                .forEach(e -> tuple.appendChild(e.accept(this)));
            return tuple;
        }

        @Override
        public Node visitNode(INode o) throws RuntimeException {
            Element node = new Element(o.getName().replaceAll("-", ":"));

            Map<String, IValue> parameters = o.asWithKeywordParameters().getParameters();

            if (parameters.containsKey("xmlns")) {
                IMap ps = (IMap) parameters.get("xmlns");
                for (IValue e : ps) {
                    IString skey = (IString) e;
                    if (skey.getValue().equals("xml")) {
                        node.attr("xmlns", ((IString) ps.get(skey)).getValue());
                    }
                    else {
                        node.attr("xmlns:" + skey.getValue(), ((IString) ps.get(skey)).getValue());
                    }
                }

                parameters.remove("xmlns");
            }

            // if there is both a src and an location attr, then "rascal-src" is the origin key, otherwise "src"
            String originKey = parameters.containsKey(SRC_ATTR) && parameters.containsKey(QUALIFIED_SRC_ATTR) ? QUALIFIED_SRC_ATTR : SRC_ATTR;

            parameters.entrySet().stream()
                .filter(v -> !dropOrigins || !v.getKey().equals(originKey))
                .forEach(e -> {
                    IValue v = e.getValue();
                    node.attr(deNormalizeAttr(e.getKey()), v.getType().isString() ? ((IString) v).getValue() : v.toString());
                });

            StreamSupport.stream(o.spliterator(), false)
                .forEach(e -> node.appendChild(e.accept(this)));

            return node;
        }

        @Override
        public Node visitConstructor(IConstructor o) throws RuntimeException {
            return visitNode(o);
        }

        @Override
        public Node visitInteger(IInteger o) throws RuntimeException {
            Element integer = new Element("integer");
            integer.attr("val", o.toString());
            return integer;
        }

        @Override
        public Node visitMap(IMap o) throws RuntimeException {
            Element map = new Element("map");

            for (Entry<IValue, IValue> e : ((Iterable<Entry<IValue,IValue>>) () -> o.entryIterator())) {
                map.appendChild(
                    new Element("entry")
                        .appendChild(e.getKey().accept(this))
                        .appendChild(e.getValue().accept(this))
                );
            }
            
            return map;
        }

        @Override
        public Node visitBoolean(IBool boolValue) throws RuntimeException {
            Element bool = new Element("bool");
            bool.attr("val", boolValue.toString());
            return bool;
        }

        @Override
        public Node visitExternal(IExternalValue externalValue) throws RuntimeException {
            return new TextNode(externalValue.toString());
        }

        @Override
        public Node visitDateTime(IDateTime o) throws RuntimeException {
            Element datetime = new Element("datetime");
            datetime.attr("val", o.toString());
            return datetime;
        }
    }

    private static String normalizeAttr(String attr) {
        return attr.replaceAll(":", "-");
    }

    private static String deNormalizeAttr(String attr) {
        return attr.replaceAll("-", ":");
    }
}