PDF logical structure reader in Swift

//---------------------------------------------------------------------------------------
// Copyright (c) 2001-2019 by PDFTron Systems Inc. All Rights Reserved.
// Consult legal.txt regarding legal and license information.
//---------------------------------------------------------------------------------------

import PDFNet
import Foundation

//---------------------------------------------------------------------------------------
// This sample explores the structure and content of a tagged PDF document and dumps
// the structure information to the console window.
//
// In tagged PDF documents StructTree acts as a central repository for information
// related to a PDF document's logical structure. The tree consists of StructElement-s
// and ContentItem-s which are leaf nodes of the structure tree.
//
// The sample can be extended to access and extract the marked-content elements such
// as text and images.
//---------------------------------------------------------------------------------------

func PrintIndent(_ indent: Int) -> String {
    var str = "\n"
    for _ in 0..<indent {
        str += "  "
    }
    return str
}

// Used in code snippet 1.
func ProcessStructElement(element: PTSElement, indent: Int) -> String {
    if !element.isValid() {
        return ""
    }
    
    // Print out the type and title info, if any.
    var result = ("\(PrintIndent(indent))Type: \(String(describing: element.getType()))")
    let nestedIndent = indent + 1
    if element.hasTitle() {
        result = result + (". Title: \(String(describing: element.getTitle()))")
    }
    
    let num = element.getNumKids()
    for i in 0..<num {
        // Check if the kid is a leaf node (i.e. it is a ContentItem).
        if element.isContentItem(i) {
            let cont: PTContentItem = element.getAsContentItem(i)
            let type: PTContentItemType = cont.getType()
            
            let page: PTPage = cont.getPage()
            result += ("\(PrintIndent(nestedIndent))Content Item. Part of page #\(page.getIndex())\(PrintIndent(nestedIndent))")
            switch type {
            case e_ptMCID, e_ptMCR:
                result += ("MCID: \(cont.getMCID())")
            case e_ptOBJR:
                result += ("OBJR ")
                if let ref_obj = cont.getRefObj() {
                    result += ("- Referenced Object#: \(ref_obj.getNum())")
                }
            default:
                break
            }
        }
        else {
            // the kid is another StructElement node.
            result = result + (ProcessStructElement(element: element.getAsStructElem(i), indent: nestedIndent))
        }
    }
    return result
}

// Used in code snippet 2.
func ProcessLogicalStructureTestElements(reader: PTElementReader) -> String {
    var result = ""
    while let element = reader.next() { // Read page contents
        // In this sample we process only paths & text, but the code can be
        // extended to handle any element type.
        let type: PTElementType = element.getType()
        if type == e_ptpath || type == e_pttext_obj || type == e_ptpath {
            switch type {
            case e_ptpath:  // Process path ...
                result = result + ("\nPATH: ")
            case e_pttext_obj:  // Process text ...
                result = result + ("\nTEXT: \(String(describing: element.getTextString()))")
            case e_ptform:  // Process form XObjects
                result = result + ("\nFORM XObject:")
                //reader.FormBegin();
                //ProcessLogicalStructureTestElements(reader);
                //reader.End();
            default:
                break
            }
            
            // Check if the element is associated with any structural element.
            // Content items are leaf nodes of the structure tree.
            let struct_parent: PTSElement = element.getParentStructElement()
            if struct_parent.isValid() {
                // Print out the parent structural element's type, title, and object number.
                result = result + (" Type: \(String(describing: struct_parent.getType())), MCID: \(element.getStructMCID())")
                if struct_parent.hasTitle() {
                    result = result + (". Title: \(String(describing: struct_parent.getTitle()))")
                }
                result = result + (", Obj#: \(struct_parent.getSDFObj().getNum())")
            }
        }
    }
    return result
}


// Used in code snippet 3.
//typedef map<int, string> MCIDPageMap;
//var MCIDPageMap = [AnyHashable: Any]()
//var MCIDDocMap = [AnyHashable: Any]()
//typedef map<int, MCIDPageMap> MCIDDocMap;

// Used in code snippet 3.
func ProcessLogicalStructureTestElements2(reader: PTElementReader, mcid_page_map: NSMutableDictionary) {
    while let element = reader.next() {    // Read page contents
        // In this sample we process only text, but the code can be extended
        // to handle paths, images, or any other Element type.
        let mcid = element.getStructMCID()
        if mcid >= 0 && element.getType() == e_pttext_obj {
            let val = element.getTextString()
            let key = mcid
            if let str = mcid_page_map[key] as? String {
                mcid_page_map[key] = str + (val ?? "")
            } else {
                mcid_page_map[key] = val ?? ""
            }
        }
    }
}

// Used in code snippet 3.
func ProcessStructElement2(element: PTSElement, mcid_doc_map: NSMutableDictionary, indent: Int) -> String {
    if !element.isValid() {
        return ""
    }
    var result = ""
    // Print out the type and title info, if any.
    result += (PrintIndent(indent))
    result += ("<\(String(describing: element.getType()))")
    if element.hasTitle() {
        result += (" title=\"\(String(describing: element.getTitle()))\"")
    }
    result += (">")
    
    let num = element.getNumKids()
    for i in 0..<num {
        if element.isContentItem(i) {
            let cont: PTContentItem = element.getAsContentItem(i)
            if cont.getType() == e_ptMCID {
                let page_num = cont.getPage().getIndex()
                let key = page_num
                if let mcid_page_map = mcid_doc_map[key] as? NSMutableDictionary {
                    let key2 = cont.getMCID()
                    if let str = mcid_page_map[key2] as? String {
                        result += (str)
                    }
                }
            }
        }
        else {  // the kid is another StructElement node.
            result += (ProcessStructElement2(element: element.getAsStructElem(i), mcid_doc_map: mcid_doc_map, indent: indent + 1))
        }
    }
    
    result += (PrintIndent(indent))
    result += ("</\(String(describing: element.getType()))>")
    return result
}

func runLogicalStructureTest() -> Int {
    return autoreleasepool {
        var ret: Int = 0
        
        
        do {
            // Extract logical structure from a PDF document
            try PTPDFNet.catchException {
                let doc: PTPDFDoc = PTPDFDoc(filepath: Bundle.main.path(forResource: "tagged", ofType: "pdf"))
                doc.initSecurityHandler()
                
                print("____________________________________________________________")
                print("Sample 1 - Traverse logical structure tree...")
                do {
                    let tree: PTSTree = doc.getStructTree()
                    if tree.isValid() {
                        print("Document has a StructTree root.")
                
                        for i in 0..<tree.getNumKids() {
                            // Recursively get structure info for all child elements.
                            print("\(ProcessStructElement(element: tree.getKid(i), indent: 0))")
                        }
                    }
                    else {
                        print("This document does not contain any logical structure.")
                    }
                }
                print("Done 1.")

                print("____________________________________________________________")
                print("Sample 2 - Get parent logical structure elements from")
                print("layout elements.")
                do {
                    let reader: PTElementReader = PTElementReader()
                    let itr: PTPageIterator = doc.getPageIterator(1)
                    while itr.hasNext() {
                        reader.begin(itr.current())
                        print("\(ProcessLogicalStructureTestElements(reader: reader))")
                        reader.end()
                        itr.next()
                    }
                }
                print("Done 2.")
                
                print("____________________________________________________________")
                print("Sample 3 - 'XML style' extraction of PDF logical structure and page content.")
                do {
                    let mcid_doc_map = NSMutableDictionary()
                    let reader: PTElementReader = PTElementReader()
                    let itr: PTPageIterator = doc.getPageIterator(1)
                    while itr.hasNext() {
                        reader.begin(itr.current())
                        let arr = NSMutableDictionary()
                        ProcessLogicalStructureTestElements2(reader: reader, mcid_page_map: arr)
                        let key = itr.current().getIndex()
                        mcid_doc_map[key] = arr
                        reader.end()
                        itr.next()
                    }
                    
                    let tree: PTSTree = doc.getStructTree()
                    if tree.isValid() {
                        for i in 0..<tree.getNumKids() {
                            print("\(ProcessStructElement2(element: tree.getKid(i), mcid_doc_map: mcid_doc_map, indent: 0))")
                        }
                    }
                }
                print("Done 3.")
            }
        } catch let e as NSError {
            print("Caught PDFNet exception: \(e)")
            ret = 1
        }
        return ret
    }
}