# Copyright (C) 2011 Apple Inc. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. 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.
#
# THIS SOFTWARE IS PROVIDED BY APPLE INC. AND ITS 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 APPLE INC. OR ITS 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.

require "config"

#
# Base utility types for the AST.
#

# Valid methods for Node:
#
# node.children -> Returns an array of immediate children.
#
# node.descendents -> Returns an array of all strict descendants (children
#     and children of children, transitively).
#
# node.flatten -> Returns an array containing the strict descendants and
#     the node itself.
#
# node.filter(type) -> Returns an array containing those elements in
#     node.flatten that are of the given type (is_a? type returns true).
#
# node.mapChildren{|v| ...} -> Returns a new node with all children
#     replaced according to the given block.
#
# Examples:
#
# node.filter(Setting).uniq -> Returns all of the settings that the AST's
#     IfThenElse blocks depend on.
#
# node.filter(StructOffset).uniq -> Returns all of the structure offsets
#     that the AST depends on.

class Node
    attr_reader :codeOrigin
    
    def initialize(codeOrigin)
        @codeOrigin = codeOrigin
    end
    
    def codeOriginString
        @codeOrigin.to_s
    end
    
    def descendants
        children.collect{|v| v.flatten}.flatten
    end
    
    def flatten
        [self] + descendants
    end
    
    def filter(type)
        flatten.select{|v| v.is_a? type}
    end
end

class NoChildren < Node
    def initialize(codeOrigin)
        super(codeOrigin)
    end
    
    def children
        []
    end
    
    def mapChildren
        self
    end
end

class StructOffsetKey
    attr_reader :struct, :field
    
    def initialize(struct, field)
        @struct = struct
        @field = field
    end
    
    def hash
        @struct.hash + @field.hash * 3
    end
    
    def eql?(other)
        @struct == other.struct and @field == other.field
    end
end

#
# AST nodes.
#

class StructOffset < NoChildren
    attr_reader :struct, :field
    
    def initialize(codeOrigin, struct, field)
        super(codeOrigin)
        @struct = struct
        @field = field
    end
    
    @@mapping = {}
    
    def self.forField(codeOrigin, struct, field)
        key = StructOffsetKey.new(struct, field)
        
        unless @@mapping[key]
            @@mapping[key] = StructOffset.new(codeOrigin, struct, field)
        end
        @@mapping[key]
    end
    
    def dump
        "#{struct}::#{field}"
    end
    
    def <=>(other)
        if @struct != other.struct
            return @struct <=> other.struct
        end
        @field <=> other.field
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def register?
        false
    end
end

class Sizeof < NoChildren
    attr_reader :struct
    
    def initialize(codeOrigin, struct)
        super(codeOrigin)
        @struct = struct
    end
    
    @@mapping = {}
    
    def self.forName(codeOrigin, struct)
        unless @@mapping[struct]
            @@mapping[struct] = Sizeof.new(codeOrigin, struct)
        end
        @@mapping[struct]
    end
    
    def dump
        "sizeof #{@struct}"
    end
    
    def <=>(other)
        @struct <=> other.struct
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def register?
        false
    end
end

class Immediate < NoChildren
    attr_reader :value
    
    def initialize(codeOrigin, value)
        super(codeOrigin)
        @value = value
        raise "Bad immediate value #{value.inspect} at #{codeOriginString}" unless value.is_a? Integer
    end
    
    def dump
        "#{value}"
    end
    
    def ==(other)
        other.is_a? Immediate and other.value == @value
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        true
    end
        
    def register?
        false
    end
end

class AddImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        AddImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} + #{right.dump})"
    end
    
    def value
        "#{left.value} + #{right.value}"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        true
    end
    
    def register?
        false
    end
end

class SubImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        SubImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} - #{right.dump})"
    end
    
    def value
        "#{left.value} - #{right.value}"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        true
    end
    
    def register?
        false
    end
end

class MulImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        MulImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} * #{right.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class NegImmediate < Node
    attr_reader :child
    
    def initialize(codeOrigin, child)
        super(codeOrigin)
        @child = child
    end
    
    def children
        [@child]
    end
    
    def mapChildren
        NegImmediate.new(codeOrigin, (yield @child))
    end
    
    def dump
        "(-#{@child.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class OrImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        OrImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} | #{right.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class AndImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        AndImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} & #{right.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class XorImmediates < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        XorImmediates.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} ^ #{right.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class BitnotImmediate < Node
    attr_reader :child
    
    def initialize(codeOrigin, child)
        super(codeOrigin)
        @child = child
    end
    
    def children
        [@child]
    end
    
    def mapChildren
        BitnotImmediate.new(codeOrigin, (yield @child))
    end
    
    def dump
        "(~#{@child.dump})"
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        true
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class StringLiteral < NoChildren
    attr_reader :value
    
    def initialize(codeOrigin, value)
        super(codeOrigin)
        @value = value[1..-2]
        raise "Bad string literal #{value.inspect} at #{codeOriginString}" unless value.is_a? String
    end
    
    def dump
        "#{value}"
    end
    
    def ==(other)
        other.is_a? StringLiteral and other.value == @value
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        false
    end
        
    def register?
        false
    end
end

class RegisterID < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
    end
    
    @@mapping = {}
    
    def self.forName(codeOrigin, name)
        unless @@mapping[name]
            @@mapping[name] = RegisterID.new(codeOrigin, name)
        end
        @@mapping[name]
    end
    
    def dump
        name
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def register?
        true
    end
end

class FPRegisterID < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
    end
    
    @@mapping = {}
    
    def self.forName(codeOrigin, name)
        unless @@mapping[name]
            @@mapping[name] = FPRegisterID.new(codeOrigin, name)
        end
        @@mapping[name]
    end
    
    def dump
        name
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        true
    end
end

class SpecialRegister < NoChildren
    def initialize(name)
        @name = name
    end
    
    def address?
        false
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        true
    end
end

class Variable < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
    end
    
    @@mapping = {}
    
    def self.forName(codeOrigin, name)
        unless @@mapping[name]
            @@mapping[name] = Variable.new(codeOrigin, name)
        end
        @@mapping[name]
    end
    
    def dump
        name
    end
    
    def inspect
        "<variable #{name} at #{codeOriginString}>"
    end
end

class Address < Node
    attr_reader :base, :offset
    
    def initialize(codeOrigin, base, offset)
        super(codeOrigin)
        @base = base
        @offset = offset
        raise "Bad base for address #{base.inspect} at #{codeOriginString}" unless base.is_a? Variable or base.register?
        raise "Bad offset for address #{offset.inspect} at #{codeOriginString}" unless offset.is_a? Variable or offset.immediate?
    end
    
    def withOffset(extraOffset)
        Address.new(codeOrigin, @base, Immediate.new(codeOrigin, @offset.value + extraOffset))
    end
    
    def children
        [@base, @offset]
    end
    
    def mapChildren
        Address.new(codeOrigin, (yield @base), (yield @offset))
    end
    
    def dump
        "#{offset.dump}[#{base.dump}]"
    end
    
    def address?
        true
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        true
    end
    
    def register?
        false
    end
end

class BaseIndex < Node
    attr_reader :base, :index, :scale, :offset
    
    def initialize(codeOrigin, base, index, scale, offset)
        super(codeOrigin)
        @base = base
        @index = index
        @scale = scale
        raise unless [1, 2, 4, 8].member? @scale
        @offset = offset
    end
    
    def scaleShift
        case scale
        when 1
            0
        when 2
            1
        when 4
            2
        when 8
            3
        else
            raise "Bad scale at #{codeOriginString}"
        end
    end
    
    def withOffset(extraOffset)
        BaseIndex.new(codeOrigin, @base, @index, @scale, Immediate.new(codeOrigin, @offset.value + extraOffset))
    end
    
    def children
        [@base, @index, @offset]
    end
    
    def mapChildren
        BaseIndex.new(codeOrigin, (yield @base), (yield @index), @scale, (yield @offset))
    end
    
    def dump
        "#{offset.dump}[#{base.dump}, #{index.dump}, #{scale}]"
    end
    
    def address?
        true
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        false
    end
    
    def register?
        false
    end
end

class AbsoluteAddress < NoChildren
    attr_reader :address
    
    def initialize(codeOrigin, address)
        super(codeOrigin)
        @address = address
    end
    
    def withOffset(extraOffset)
        AbsoluteAddress.new(codeOrigin, Immediate.new(codeOrigin, @address.value + extraOffset))
    end
    
    def dump
        "#{address.dump}[]"
    end
    
    def address?
        true
    end
    
    def label?
        false
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        true
    end
    
    def register?
        false
    end
end

class Instruction < Node
    attr_reader :opcode, :operands, :annotation
    
    def initialize(codeOrigin, opcode, operands, annotation=nil)
        super(codeOrigin)
        @opcode = opcode
        @operands = operands
        @annotation = annotation
    end
    
    def children
        operands
    end
    
    def mapChildren(&proc)
        Instruction.new(codeOrigin, @opcode, @operands.map(&proc), @annotation)
    end
    
    def dump
        "\t" + opcode.to_s + " " + operands.collect{|v| v.dump}.join(", ")
    end

    def lowerDefault
        case opcode
        when "localAnnotation"
            $asm.putLocalAnnotation
        when "globalAnnotation"
            $asm.putGlobalAnnotation
        when "emit"
          $asm.puts "#{operands[0].dump}"
        else
            raise "Unhandled opcode #{opcode} at #{codeOriginString}"
        end
    end
end

class Error < NoChildren
    def initialize(codeOrigin)
        super(codeOrigin)
    end
    
    def dump
        "\terror"
    end
end

class ConstDecl < Node
    attr_reader :variable, :value
    
    def initialize(codeOrigin, variable, value)
        super(codeOrigin)
        @variable = variable
        @value = value
    end
    
    def children
        [@variable, @value]
    end
    
    def mapChildren
        ConstDecl.new(codeOrigin, (yield @variable), (yield @value))
    end
    
    def dump
        "const #{@variable.dump} = #{@value.dump}"
    end
end

$labelMapping = {}
$referencedExternLabels = Array.new

class Label < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
        @extern = true
        @global = false
    end
    
    def self.forName(codeOrigin, name, definedInFile = false)
        if $labelMapping[name]
            raise "Label name collision: #{name}" unless $labelMapping[name].is_a? Label
        else
            $labelMapping[name] = Label.new(codeOrigin, name)
        end
        if definedInFile
            $labelMapping[name].clearExtern()
        end
        $labelMapping[name]
    end

    def self.setAsGlobal(codeOrigin, name)
        if $labelMapping[name]
            label = $labelMapping[name]
            raise "Label: #{name} declared global multiple times" unless not label.global?
            label.setGlobal()
        else
            newLabel = Label.new(codeOrigin, name)
            newLabel.setGlobal()
            $labelMapping[name] = newLabel
        end
    end

    def self.resetReferenced
        $referencedExternLabels = Array.new
    end

    def self.forReferencedExtern()
        $referencedExternLabels.each {
            | label |
            yield "#{label.name}"
        }
    end

    def clearExtern
        @extern = false
    end

    def extern?
        @extern
    end

    def setGlobal
        @global = true
    end

    def global?
        @global
    end

    def dump
        "#{name}:"
    end
end

class LocalLabel < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
    end

    @@uniqueNameCounter = 0
    
    def self.forName(codeOrigin, name)
        if $labelMapping[name]
            raise "Label name collision: #{name}" unless $labelMapping[name].is_a? LocalLabel
        else
            $labelMapping[name] = LocalLabel.new(codeOrigin, name)
        end
        $labelMapping[name]
    end
    
    def self.unique(comment)
        newName = "_#{comment}"
        if $labelMapping[newName]
            while $labelMapping[newName = "_#{@@uniqueNameCounter}_#{comment}"]
                @@uniqueNameCounter += 1
            end
        end
        forName(nil, newName)
    end
    
    def cleanName
        if name =~ /^\./
            "_" + name[1..-1]
        else
            name
        end
    end
    
    def dump
        "#{name}:"
    end
end

class LabelReference < Node
    attr_reader :label
    
    def initialize(codeOrigin, label)
        super(codeOrigin)
        @label = label
    end
    
    def children
        [@label]
    end
    
    def mapChildren
        LabelReference.new(codeOrigin, (yield @label))
    end
    
    def name
        label.name
    end
    
    def extern?
        $labelMapping[name].is_a? Label and $labelMapping[name].extern?
    end

    def used
        if !$referencedExternLabels.include?(@label) and extern?
            $referencedExternLabels.push(@label)
        end
    end

    def dump
        label.name
    end
    
    def value
        asmLabel()
    end

    def address?
        false
    end
    
    def label?
        true
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        true
    end
end

class LocalLabelReference < NoChildren
    attr_reader :label
    
    def initialize(codeOrigin, label)
        super(codeOrigin)
        @label = label
    end
    
    def children
        [@label]
    end
    
    def mapChildren
        LocalLabelReference.new(codeOrigin, (yield @label))
    end
    
    def name
        label.name
    end
    
    def dump
        label.name
    end

    def value
        asmLabel()
    end
    
    def address?
        false
    end
    
    def label?
        true
    end
    
    def immediate?
        false
    end
    
    def immediateOperand?
        true
    end
end

class Sequence < Node
    attr_reader :list
    
    def initialize(codeOrigin, list)
        super(codeOrigin)
        @list = list
    end
    
    def children
        list
    end
    
    def mapChildren(&proc)
        Sequence.new(codeOrigin, @list.map(&proc))
    end
    
    def dump
        list.collect{|v| v.dump}.join("\n")
    end
end

class True < NoChildren
    def initialize
        super(nil)
    end
    
    @@instance = True.new
    
    def self.instance
        @@instance
    end
    
    def value
        true
    end
    
    def dump
        "true"
    end
end

class False < NoChildren
    def initialize
        super(nil)
    end
    
    @@instance = False.new
    
    def self.instance
        @@instance
    end
    
    def value
        false
    end
    
    def dump
        "false"
    end
end

class TrueClass
    def asNode
        True.instance
    end
end

class FalseClass
    def asNode
        False.instance
    end
end

class Setting < NoChildren
    attr_reader :name
    
    def initialize(codeOrigin, name)
        super(codeOrigin)
        @name = name
    end
    
    @@mapping = {}
    
    def self.forName(codeOrigin, name)
        unless @@mapping[name]
            @@mapping[name] = Setting.new(codeOrigin, name)
        end
        @@mapping[name]
    end
    
    def dump
        name
    end
end

class And < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        And.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} and #{right.dump})"
    end
end

class Or < Node
    attr_reader :left, :right
    
    def initialize(codeOrigin, left, right)
        super(codeOrigin)
        @left = left
        @right = right
    end
    
    def children
        [@left, @right]
    end
    
    def mapChildren
        Or.new(codeOrigin, (yield @left), (yield @right))
    end
    
    def dump
        "(#{left.dump} or #{right.dump})"
    end
end

class Not < Node
    attr_reader :child
    
    def initialize(codeOrigin, child)
        super(codeOrigin)
        @child = child
    end
    
    def children
        [@child]
    end
    
    def mapChildren
        Not.new(codeOrigin, (yield @child))
    end
    
    def dump
        "(not #{child.dump})"
    end
end

class Skip < NoChildren
    def initialize(codeOrigin)
        super(codeOrigin)
    end
    
    def dump
        "\tskip"
    end
end

class IfThenElse < Node
    attr_reader :predicate, :thenCase
    attr_accessor :elseCase
    
    def initialize(codeOrigin, predicate, thenCase)
        super(codeOrigin)
        @predicate = predicate
        @thenCase = thenCase
        @elseCase = Skip.new(codeOrigin)
    end
    
    def children
        if @elseCase
            [@predicate, @thenCase, @elseCase]
        else
            [@predicate, @thenCase]
        end
    end
    
    def mapChildren
        IfThenElse.new(codeOrigin, (yield @predicate), (yield @thenCase), (yield @elseCase))
    end
    
    def dump
        "if #{predicate.dump}\n" + thenCase.dump + "\nelse\n" + elseCase.dump + "\nend"
    end
end

class Macro < Node
    attr_reader :name, :variables, :body

    def initialize(codeOrigin, name, variables, body)
        super(codeOrigin)
        @name = name
        @variables = variables
        @body = body
    end
    
    def children
        @variables + [@body]
    end
    
    def mapChildren
        Macro.new(codeOrigin, @name, @variables.map{|v| yield v}, (yield @body))
    end
    
    def dump
        "macro #{name}(" + variables.collect{|v| v.dump}.join(", ") + ")\n" + body.dump + "\nend"
    end
end

class MacroCall < Node
    attr_reader :name, :operands, :annotation
    
    def initialize(codeOrigin, name, operands, annotation)
        super(codeOrigin)
        @name = name
        @operands = operands
        raise unless @operands
        @operands.each{|v| raise unless v}
        @annotation = annotation
    end
    
    def children
        @operands
    end
    
    def mapChildren(&proc)
        MacroCall.new(codeOrigin, @name, @operands.map(&proc), @annotation)
    end
    
    def dump
        "\t#{name}(" + operands.collect{|v| v.dump}.join(", ") + ")"
    end
end