// Copyright 2017 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocore

import (
	"debug/dwarf"
	"fmt"
	"strings"

	"golang.org/x/debug/core"
)

// A Process represents the state of a Go process that core dumped.
type Process struct {
	proc *core.Process

	arenaStart core.Address
	arenaUsed  core.Address
	bitmapEnd  core.Address

	// data structure for fast object finding
	heapInfo []heapInfo

	// number of live objects
	nObj int

	goroutines []*Goroutine

	// runtime info
	rtGlobals   map[string]region
	rtConstants map[string]int64

	// A module is a loadable unit. Most Go programs have 1, programs
	// which load plugins will have more.
	modules []*module

	// address -> function mapping
	funcTab funcTab

	// map from dwarf type to *Type
	dwarfMap map[dwarf.Type]*Type

	// map from address of runtime._type to *Type
	runtimeMap map[core.Address]*Type

	// map from runtime type name to the set of *Type with that name
	// Used to find candidates to put in the runtimeMap map.
	runtimeNameMap map[string][]*Type

	// memory usage by category
	stats *Stats

	buildVersion string

	globals []*Root

	// Types of each object, indexed by object index.
	// Only initialized if FlagTypes is passed to Core.
	types []typeInfo

	// Reverse edges.
	// The reverse edges for object #i are redge[ridx[i]:ridx[i+1]].
	// A "reverse edge" for object #i is a location in memory where a pointer
	// to object #i lives.
	// Only initialized if FlagReverse is passed to Core.
	redge []core.Address
	ridx  []int64
	// Sorted list of all roots.
	// Only initialized if FlagReverse is passed to Core.
	rootIdx []*Root
}

// Process returns the core.Process used to construct this Process.
func (p *Process) Process() *core.Process {
	return p.proc
}

func (p *Process) Goroutines() []*Goroutine {
	return p.goroutines
}

// Stats returns a breakdown of the program's memory use by category.
func (p *Process) Stats() *Stats {
	return p.stats
}

// BuildVersion returns the Go version that was used to build the inferior binary.
func (p *Process) BuildVersion() string {
	return p.buildVersion
}

func (p *Process) Globals() []*Root {
	return p.globals
}

// FindFunc returns the function which contains the code at address pc, if any.
func (p *Process) FindFunc(pc core.Address) *Func {
	return p.funcTab.find(pc)
}

func (p *Process) findType(name string) *Type {
	s := p.runtimeNameMap[name]
	if len(s) == 0 {
		panic("can't find type " + name)
	}
	return s[0]
}

// A Flags indicates optional analyses for Core to compute.
type Flags uint8

const (
	// FlagTypes requests that Core compute type information for all Go objects,
	// required to use the Type function.
	// Setting this flag will require more initialization time and use more memory.
	FlagTypes Flags = 1 << iota
	// FlagReverse requests that Core compute reverse edge information,
	// required to use ForEachReversePtr.
	// Setting this flag will require more initialization time and use more memory.
	FlagReverse
)

// Core takes a loaded core file and extracts Go information from it.
// flags is a bitmask of data that should be extracted from the core.
func Core(proc *core.Process, flags Flags) (p *Process, err error) {
	// Make sure we have DWARF info.
	if _, err := proc.DWARF(); err != nil {
		return nil, err
	}

	// Guard against failures of proc.Read* routines.
	/*
		defer func() {
			e := recover()
			if e == nil {
				return
			}
			p = nil
			if x, ok := e.(error); ok {
				err = x
				return
			}
			panic(e) // Not an error, re-panic it.
		}()
	*/

	p = &Process{
		proc:       proc,
		runtimeMap: map[core.Address]*Type{},
		dwarfMap:   map[dwarf.Type]*Type{},
	}

	// Initialize everything that just depends on DWARF.
	p.readDWARFTypes()
	p.readRuntimeConstants()
	p.readGlobals()

	// Find runtime globals we care about. Initialize regions for them.
	p.rtGlobals = map[string]region{}
	for _, g := range p.globals {
		if strings.HasPrefix(g.Name, "runtime.") {
			p.rtGlobals[g.Name[8:]] = region{p: p, a: g.Addr, typ: g.Type}
		}
	}

	// Read all the data that depend on runtime globals.
	p.buildVersion = p.rtGlobals["buildVersion"].String()
	p.readModules()
	p.readSpans()
	p.readGs()
	p.readStackVars() // needs to be after readGs.
	p.markObjects()   // needs to be after readGlobals, readStackVars.
	if flags&FlagTypes != 0 {
		p.typeHeap() // needs to be after markObjects.
	}
	if flags&FlagReverse != 0 {
		p.reverseEdges() // needs to be after markObjects.
	}

	return p, nil
}

func (p *Process) readSpans() {
	mheap := p.rtGlobals["mheap_"]

	spanTableStart := mheap.Field("spans").SlicePtr().Address()
	spanTableEnd := spanTableStart.Add(mheap.Field("spans").SliceCap() * p.proc.PtrSize())
	arenaStart := core.Address(mheap.Field("arena_start").Uintptr())
	arenaUsed := core.Address(mheap.Field("arena_used").Uintptr())
	arenaEnd := core.Address(mheap.Field("arena_end").Uintptr())
	bitmapEnd := core.Address(mheap.Field("bitmap").Uintptr())
	bitmapStart := bitmapEnd.Add(-int64(mheap.Field("bitmap_mapped").Uintptr()))

	p.arenaStart = arenaStart
	p.arenaUsed = arenaUsed
	p.bitmapEnd = bitmapEnd

	var all int64
	var text int64
	var readOnly int64
	var heap int64
	var spanTable int64
	var bitmap int64
	var data int64
	var bss int64 // also includes mmap'd regions
	for _, m := range p.proc.Mappings() {
		size := m.Size()
		all += size
		switch m.Perm() {
		case core.Read:
			readOnly += size
		case core.Read | core.Exec:
			text += size
		case core.Read | core.Write:
			if m.CopyOnWrite() {
				// Check if m.file == text's file? That could distinguish
				// data segment from mmapped file.
				data += size
				break
			}
			attribute := func(x, y core.Address, p *int64) {
				a := x.Max(m.Min())
				b := y.Min(m.Max())
				if a < b {
					*p += b.Sub(a)
					size -= b.Sub(a)
				}
			}
			attribute(spanTableStart, spanTableEnd, &spanTable)
			attribute(arenaStart, arenaEnd, &heap)
			attribute(bitmapStart, bitmapEnd, &bitmap)
			// Any other anonymous mapping is bss.
			// TODO: how to distinguish original bss from anonymous mmap?
			bss += size
		default:
			panic("weird mapping " + m.Perm().String())
		}
	}
	pageSize := p.rtConstants["_PageSize"]

	// Span types
	spanInUse := uint8(p.rtConstants["_MSpanInUse"])
	spanManual := uint8(p.rtConstants["_MSpanManual"])
	spanDead := uint8(p.rtConstants["_MSpanDead"])
	spanFree := uint8(p.rtConstants["_MSpanFree"])

	// Process spans.
	if pageSize%heapInfoSize != 0 {
		panic(fmt.Sprintf("page size not a multiple of %d", heapInfoSize))
	}
	p.heapInfo = make([]heapInfo, (p.arenaUsed-p.arenaStart)/heapInfoSize)
	allspans := mheap.Field("allspans")
	var allSpanSize int64
	var freeSpanSize int64
	var manualSpanSize int64
	var inUseSpanSize int64
	var allocSize int64
	var freeSize int64
	var spanRoundSize int64
	var manualAllocSize int64
	var manualFreeSize int64
	n := allspans.SliceLen()
	for i := int64(0); i < n; i++ {
		s := allspans.SliceIndex(i).Deref()
		min := core.Address(s.Field("startAddr").Uintptr())
		elemSize := int64(s.Field("elemsize").Uintptr())
		nPages := int64(s.Field("npages").Uintptr())
		spanSize := nPages * pageSize
		max := min.Add(spanSize)
		allSpanSize += spanSize
		switch s.Field("state").Uint8() {
		case spanInUse:
			inUseSpanSize += spanSize
			n := int64(s.Field("nelems").Uintptr())
			// An object is allocated if it is marked as
			// allocated or it is below freeindex.
			x := s.Field("allocBits").Address()
			alloc := make([]bool, n)
			for i := int64(0); i < n; i++ {
				alloc[i] = p.proc.ReadUint8(x.Add(i/8))>>uint(i%8)&1 != 0
			}
			k := int64(s.Field("freeindex").Uintptr())
			for i := int64(0); i < k; i++ {
				alloc[i] = true
			}
			for i := int64(0); i < n; i++ {
				if alloc[i] {
					allocSize += elemSize
				} else {
					freeSize += elemSize
				}
			}
			spanRoundSize += spanSize - n*elemSize
			for a := min; a < max; a += heapInfoSize {
				p.heapInfo[(a.Sub(p.arenaStart))/heapInfoSize] = heapInfo{base: min, size: elemSize, firstIdx: -1}
			}

			// Process special records.
			for sp := s.Field("specials"); sp.Address() != 0; sp = sp.Field("next") {
				sp = sp.Deref() // *special to special
				if sp.Field("kind").Uint8() != uint8(p.rtConstants["_KindSpecialFinalizer"]) {
					// All other specials (just profile records) can't point into the heap.
					continue
				}
				obj := min.Add(int64(sp.Field("offset").Uint16()))
				p.globals = append(p.globals,
					&Root{
						Name:  fmt.Sprintf("finalizer for %x", obj),
						Addr:  sp.a,
						Type:  p.findType("runtime.specialfinalizer"),
						Frame: nil,
					})
				// TODO: these aren't really "globals", as they
				// are kept alive by the object they reference being alive.
				// But we have no way of adding edges from an object to
				// the corresponding finalizer data, so we punt on that thorny
				// issue for now.
			}
		case spanFree:
			freeSpanSize += spanSize
		case spanDead:
			// These are just deallocated span descriptors. They use no heap.
		case spanManual:
			manualSpanSize += spanSize
			manualAllocSize += spanSize
			for x := core.Address(s.Field("manualFreeList").Cast("uintptr").Uintptr()); x != 0; x = p.proc.ReadPtr(x) {
				manualAllocSize -= elemSize
				manualFreeSize += elemSize
			}
		}
	}

	p.stats = &Stats{"all", all, []*Stats{
		&Stats{"text", text, nil},
		&Stats{"readonly", readOnly, nil},
		&Stats{"data", data, nil},
		&Stats{"bss", bss, nil},
		&Stats{"heap", heap, []*Stats{
			&Stats{"in use spans", inUseSpanSize, []*Stats{
				&Stats{"alloc", allocSize, nil},
				&Stats{"free", freeSize, nil},
				&Stats{"round", spanRoundSize, nil},
			}},
			&Stats{"manual spans", manualSpanSize, []*Stats{
				&Stats{"alloc", manualAllocSize, nil},
				&Stats{"free", manualFreeSize, nil},
			}},
			&Stats{"free spans", freeSpanSize, nil},
		}},
		&Stats{"ptr bitmap", bitmap, nil},
		&Stats{"span table", spanTable, nil},
	}}

	var check func(*Stats)
	check = func(s *Stats) {
		if len(s.Children) == 0 {
			return
		}
		var sum int64
		for _, c := range s.Children {
			sum += c.Size
		}
		if sum != s.Size {
			panic(fmt.Sprintf("check failed for %s: %d vs %d", s.Name, s.Size, sum))
		}
		for _, c := range s.Children {
			check(c)
		}
	}
	check(p.stats)
}

func (p *Process) readGs() {
	// TODO: figure out how to "flush" running Gs.
	allgs := p.rtGlobals["allgs"]
	n := allgs.SliceLen()
	for i := int64(0); i < n; i++ {
		r := allgs.SliceIndex(i).Deref()
		g := p.readG(r)
		if g == nil {
			continue
		}
		p.goroutines = append(p.goroutines, g)
	}
}

func (p *Process) readG(r region) *Goroutine {
	g := &Goroutine{r: r}
	stk := r.Field("stack")
	g.stackSize = int64(stk.Field("hi").Uintptr() - stk.Field("lo").Uintptr())

	var osT *core.Thread // os thread working on behalf of this G (if any).
	mp := r.Field("m")
	if mp.Address() != 0 {
		m := mp.Deref()
		pid := m.Field("procid").Uint64()
		// TODO check that m.curg points to g?
		for _, t := range p.proc.Threads() {
			if t.Pid() == pid {
				osT = t
			}
		}
	}
	status := r.Field("atomicstatus").Uint32()
	status &^= uint32(p.rtConstants["_Gscan"])
	var sp, pc core.Address
	switch status {
	case uint32(p.rtConstants["_Gidle"]):
		return g
	case uint32(p.rtConstants["_Grunnable"]), uint32(p.rtConstants["_Gwaiting"]):
		sched := r.Field("sched")
		sp = core.Address(sched.Field("sp").Uintptr())
		pc = core.Address(sched.Field("pc").Uintptr())
	case uint32(p.rtConstants["_Grunning"]):
		sp = osT.SP()
		pc = osT.PC()
		// TODO: back up to the calling frame?
	case uint32(p.rtConstants["_Gsyscall"]):
		sp = core.Address(r.Field("syscallsp").Uintptr())
		pc = core.Address(r.Field("syscallpc").Uintptr())
		// TODO: or should we use the osT registers?
	case uint32(p.rtConstants["_Gdead"]):
		return nil
		// TODO: copystack, others?
	default:
		// Unknown state. We can't read the frames, so just bail now.
		// TODO: make this switch complete and then panic here.
		// TODO: or just return nil?
		return g
	}
	for {
		f := p.readFrame(sp, pc)
		if f.f.name == "runtime.goexit" {
			break
		}
		if len(g.frames) > 0 {
			g.frames[len(g.frames)-1].parent = f
		}
		g.frames = append(g.frames, f)

		if f.f.name == "runtime.sigtrampgo" {
			// Continue traceback at location where the signal
			// interrupted normal execution.
			ctxt := p.proc.ReadPtr(sp.Add(16)) // 3rd arg
			//ctxt is a *ucontext
			mctxt := ctxt.Add(5 * 8)
			// mctxt is a *mcontext
			sp = p.proc.ReadPtr(mctxt.Add(15 * 8))
			pc = p.proc.ReadPtr(mctxt.Add(16 * 8))
			// TODO: totally arch-dependent!
		} else {
			sp = f.max
			pc = core.Address(p.proc.ReadUintptr(sp - 8)) // TODO:amd64 only
		}
		if pc == 0 {
			// TODO: when would this happen?
			break
		}
		if f.f.name == "runtime.systemstack" {
			// switch over to goroutine stack
			sched := r.Field("sched")
			sp = core.Address(sched.Field("sp").Uintptr())
			pc = core.Address(sched.Field("pc").Uintptr())
		}
	}
	return g
}

func (p *Process) readFrame(sp, pc core.Address) *Frame {
	f := p.funcTab.find(pc)
	if f == nil {
		panic(fmt.Errorf("  pc not found %x\n", pc))
	}
	off := pc.Sub(f.entry)
	size := f.frameSize.find(off)
	size += p.proc.PtrSize() // TODO: on amd64, the pushed return address

	frame := &Frame{f: f, pc: pc, min: sp, max: sp.Add(size)}

	// Find live ptrs in locals
	live := map[core.Address]bool{}
	if x := int(p.rtConstants["_FUNCDATA_LocalsPointerMaps"]); x < len(f.funcdata) {
		locals := region{p: p, a: f.funcdata[x], typ: p.findType("runtime.stackmap")}
		n := locals.Field("n").Int32()       // # of bitmaps
		nbit := locals.Field("nbit").Int32() // # of bits per bitmap
		idx := f.stackMap.find(off)
		if idx < 0 {
			idx = 0
		}
		if idx < int64(n) {
			bits := locals.Field("bytedata").a.Add(int64(nbit+7) / 8 * idx)
			base := frame.max.Add(-16).Add(-int64(nbit) * p.proc.PtrSize())
			// TODO: -16 for amd64. Return address and parent's frame pointer
			for i := int64(0); i < int64(nbit); i++ {
				if p.proc.ReadUint8(bits.Add(i/8))>>uint(i&7)&1 != 0 {
					live[base.Add(i*p.proc.PtrSize())] = true
				}
			}
		}
	}
	// Same for args
	if x := int(p.rtConstants["_FUNCDATA_ArgsPointerMaps"]); x < len(f.funcdata) {
		args := region{p: p, a: f.funcdata[x], typ: p.findType("runtime.stackmap")}
		n := args.Field("n").Int32()       // # of bitmaps
		nbit := args.Field("nbit").Int32() // # of bits per bitmap
		idx := f.stackMap.find(off)
		if idx < 0 {
			idx = 0
		}
		if idx < int64(n) {
			bits := args.Field("bytedata").a.Add(int64(nbit+7) / 8 * idx)
			base := frame.max
			// TODO: add to base for LR archs.
			for i := int64(0); i < int64(nbit); i++ {
				if p.proc.ReadUint8(bits.Add(i/8))>>uint(i&7)&1 != 0 {
					live[base.Add(i*p.proc.PtrSize())] = true
				}
			}
		}
	}
	frame.Live = live

	return frame
}

// A Stats struct is the node of a tree representing the entire memory
// usage of the Go program. Children of a node break its usage down
// by category.
// We maintain the invariant that, if there are children,
// Size == sum(c.Size for c in Children).
type Stats struct {
	Name     string
	Size     int64
	Children []*Stats
}

func (s *Stats) Child(name string) *Stats {
	for _, c := range s.Children {
		if c.Name == name {
			return c
		}
	}
	return nil
}