<commit_before>package parser

import (
	"fmt"

	"github.com/nitrogen-lang/nitrogen/src/ast"
	"github.com/nitrogen-lang/nitrogen/src/lexer"
	"github.com/nitrogen-lang/nitrogen/src/token"
)

const (
	priLowest      int = iota
	priCompare         // and, or
	priEquals          // ==
	priLessGreater     // > or <
	priSum             // +, -
	priProduct         // *, /
	priPrefix          // -x or !x
	priCall            // myFunction(x)
	priIndex           // array[index]
	priAssign
)

var precedences = map[token.TokenType]int{
	token.LAnd:          priCompare,
	token.LOr:           priCompare,
	token.Equal:         priEquals,
	token.NotEqual:      priEquals,
	token.LessThanEq:    priEquals,
	token.GreaterThanEq: priEquals,
	token.LessThan:      priLessGreater,
	token.GreaterThan:   priLessGreater,
	token.Plus:          priSum,
	token.Dash:          priSum,
	token.BitwiseOr:     priSum,
	token.Carrot:        priSum,
	token.Slash:         priProduct,
	token.Asterisk:      priProduct,
	token.Modulo:        priProduct,
	token.ShiftLeft:     priProduct,
	token.ShiftRight:    priProduct,
	token.BitwiseAnd:    priProduct,
	token.BitwiseAndNot: priProduct,
	token.LParen:        priCall,
	token.LSquare:       priIndex,
	token.Dot:           priIndex,
	token.Assign:        priAssign,
	token.PlusAssign:    priAssign,
	token.MinusAssign:   priAssign,
	token.TimesAssign:   priAssign,
	token.SlashAssign:   priAssign,
	token.ModAssign:     priAssign,
}

type (
	prefixParseFn func() ast.Expression
	infixParseFn  func(ast.Expression) ast.Node
	Settings      struct {
		Debug bool
	}
)

type Parser struct {
	l        *lexer.Lexer
	errors   []string
	settings *Settings

	lastToken token.Token
	curToken  token.Token
	peekToken token.Token

	insertedTokens []token.Token

	prefixParseFns map[token.TokenType]prefixParseFn
	infixParseFns  map[token.TokenType]infixParseFn
}

func New(l *lexer.Lexer, settings *Settings) *Parser {
	if settings == nil {
		settings = &Settings{}
	}

	p := &Parser{
		l:              l,
		settings:       settings,
		errors:         []string{},
		insertedTokens: make([]token.Token, 0, 5),
	}

	// Prefix parsing functions
	p.prefixParseFns = make(map[token.TokenType]prefixParseFn)
	p.registerPrefix(token.Identifier, p.parseIdentifier)
	p.registerPrefix(token.Integer, p.parseIntegerLiteral)
	p.registerPrefix(token.Float, p.parseFloatLiteral)
	p.registerPrefix(token.Nil, p.parseNullLiteral)
	p.registerPrefix(token.String, p.parseStringLiteral)
	p.registerPrefix(token.True, p.parseBoolean)
	p.registerPrefix(token.False, p.parseBoolean)
	p.registerPrefix(token.LSquare, p.parseArrayLiteral)
	p.registerPrefix(token.LBrace, p.parseHashLiteral)
	p.registerPrefix(token.Function, p.parseFunctionLiteral)
	p.registerPrefix(token.Bang, p.parsePrefixExpression)
	p.registerPrefix(token.Dash, p.parsePrefixExpression)
	p.registerPrefix(token.LParen, p.parseGroupedExpression)
	p.registerPrefix(token.If, p.parseIfExpression)
	p.registerPrefix(token.Try, p.parseTryCatch)
	p.registerPrefix(token.Class, p.parseClassLiteral)
	p.registerPrefix(token.New, p.parseMakeExpression)

	// Infix parsing functions
	p.infixParseFns = make(map[token.TokenType]infixParseFn)
	p.registerInfix(token.Plus, p.parseInfixExpression)
	p.registerInfix(token.Dash, p.parseInfixExpression)
	p.registerInfix(token.Slash, p.parseInfixExpression)
	p.registerInfix(token.Asterisk, p.parseInfixExpression)
	p.registerInfix(token.Modulo, p.parseInfixExpression)
	p.registerInfix(token.Equal, p.parseInfixExpression)
	p.registerInfix(token.NotEqual, p.parseInfixExpression)
	p.registerInfix(token.LessThanEq, p.parseInfixExpression)
	p.registerInfix(token.GreaterThanEq, p.parseInfixExpression)
	p.registerInfix(token.LessThan, p.parseInfixExpression)
	p.registerInfix(token.GreaterThan, p.parseInfixExpression)
	p.registerInfix(token.LAnd, p.parseCompareExpression)
	p.registerInfix(token.LOr, p.parseCompareExpression)
	p.registerInfix(token.LParen, p.parseCallExpression)
	p.registerInfix(token.LSquare, p.parseIndexExpression)
	p.registerInfix(token.Dot, p.parseAttributeExpression)
	p.registerInfix(token.Assign, p.parseAssignmentStatement)
	p.registerInfix(token.PlusAssign, p.parseCompoundAssign)
	p.registerInfix(token.MinusAssign, p.parseCompoundAssign)
	p.registerInfix(token.TimesAssign, p.parseCompoundAssign)
	p.registerInfix(token.SlashAssign, p.parseCompoundAssign)
	p.registerInfix(token.ModAssign, p.parseCompoundAssign)
	p.registerInfix(token.ShiftLeft, p.parseInfixExpression)
	p.registerInfix(token.ShiftRight, p.parseInfixExpression)
	p.registerInfix(token.BitwiseAnd, p.parseInfixExpression)
	p.registerInfix(token.BitwiseAndNot, p.parseInfixExpression)
	p.registerInfix(token.BitwiseOr, p.parseInfixExpression)
	p.registerInfix(token.Carrot, p.parseInfixExpression)

	// Read the first two tokens to populate curToken and peekToken
	p.nextToken()
	p.nextToken()

	return p
}

func (p *Parser) Errors() []string {
	return p.errors
}

func (p *Parser) registerPrefix(tt token.TokenType, fn prefixParseFn) {
	p.prefixParseFns[tt] = fn
}

func (p *Parser) registerInfix(tt token.TokenType, fn infixParseFn) {
	p.infixParseFns[tt] = fn
}

func (p *Parser) addError(format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	p.errors = append(p.errors, msg)
}

func (p *Parser) addErrorWithPos(format string, args ...interface{}) {
	if len(args) > 0 {
		args = append([]interface{}{p.curToken.Pos.Line, p.curToken.Pos.Col}, args...)
	} else {
		args = []interface{}{p.curToken.Pos.Line, p.curToken.Pos.Col}
	}
	msg := fmt.Sprintf("at line %d, col %d "+format, args...)
	p.errors = append(p.errors, msg)
}

func (p *Parser) peekError(t token.TokenType) {
	p.addError(
		"at line %d, col %d Incorrect next token. Expected %q, got %q",
		p.peekToken.Pos.Line,
		p.peekToken.Pos.Col,
		t.String(),
		p.peekToken.Type.String(),
	)
}

func (p *Parser) noPrefixParseFnError(t token.TokenType) {
	p.addErrorWithPos("Invalid prefix: %s", t)
}

func (p *Parser) nextToken() {
	p.advanceToken()
	for p.peekTokenIs(token.Comment) {
		p.advancePeekToken()
	}

	if p.curTokenIs(token.Illegal) {
		p.addErrorWithPos("Got illegal token: %s", p.curToken.Literal)
		p.nextToken()
	}

	if p.settings.Debug {
		fmt.Println(p.curToken)
	}
}

func (p *Parser) advanceToken() {
	p.lastToken = p.curToken
	p.curToken = p.peekToken
	p.advancePeekToken()
}

func (p *Parser) advancePeekToken() {
	if len(p.insertedTokens) > 0 {
		p.peekToken = p.insertedTokens[len(p.insertedTokens)-1]
		p.insertedTokens = p.insertedTokens[:len(p.insertedTokens)-1]
	} else {
		p.peekToken = p.l.NextToken()
	}
}

func (p *Parser) insertToken(t token.Token) {
	p.insertedTokens = append(p.insertedTokens, p.peekToken)
	p.peekToken = t
}

func (p *Parser) ParseProgram() *ast.Program {
	if p.settings.Debug {
		fmt.Println("ParseProgram")
	}
	program := &ast.Program{Filename: p.curToken.Filename}
	program.Statements = []ast.Statement{}

	for p.curToken.Type != token.EOF {
		stmt := p.parseStatement()
		if stmt != nil {
			program.Statements = append(program.Statements, stmt)
		}
		if len(p.errors) > 0 {
			return program
		}
		p.nextToken()
	}

	return program
}

func (p *Parser) parseExpression(precedence int) ast.Node {
	if p.settings.Debug {
		fmt.Println("parseExpression")
	}
	prefix := p.prefixParseFns[p.curToken.Type]
	if prefix == nil {
		p.noPrefixParseFnError(p.curToken.Type)
		return nil
	}
	var leftExp ast.Node
	leftExp = prefix()

	for !p.peekTokenIs(token.Semicolon) && precedence < p.peekPrecedence() {
		infix := p.infixParseFns[p.peekToken.Type]
		if infix == nil {
			return leftExp
		}

		p.nextToken()
		leftExp = infix(leftExp.(ast.Expression))
	}

	return leftExp
}

func (p *Parser) parseExpressionStatement() ast.Statement {
	if p.settings.Debug {
		fmt.Println("parseExpressionStatement")
	}
	var stmt ast.Statement
	curToken := p.curToken

	exp := p.parseExpression(priLowest)
	if expStmt, ok := exp.(ast.Statement); ok {
		stmt = expStmt
	} else if exp == nil {
		return nil
	} else {
		stmt = &ast.ExpressionStatement{
			Token:      curToken,
			Expression: exp.(ast.Expression),
		}
	}

	if p.peekTokenIs(token.Semicolon) {
		p.nextToken()
	}

	return stmt
}

func (p *Parser) parseBlockStatements() *ast.BlockStatement {
	if p.settings.Debug {
		fmt.Println("parseBlockStatements")
	}
	block := &ast.BlockStatement{
		Token:      p.curToken,
		Statements: []ast.Statement{},
	}
	p.nextToken()

	for !p.curTokenIs(token.RBrace) {
		stmt := p.parseStatement()
		if stmt != nil {
			block.Statements = append(block.Statements, stmt)
		}
		p.nextToken()
	}

	return block
}

func (p *Parser) parseSingleStmtBlock() *ast.BlockStatement {
	if p.settings.Debug {
		fmt.Println("parseSingleStmtBlock")
	}
	block := &ast.BlockStatement{
		Token:      p.curToken,
		Statements: make([]ast.Statement, 1),
	}
	p.nextToken()

	stmt := p.parseStatement()
	if stmt != nil {
		block.Statements[0] = stmt
	}

	return block
}

func (p *Parser) parseExpressionList(end token.TokenType) []ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseExpressionList")
	}
	list := []ast.Expression{}

	if p.peekTokenIs(end) {
		p.nextToken()
		return list
	}

	p.nextToken()
	expr := p.parseExpression(priLowest)
	if expr == nil {
		return nil
	}
	list = append(list, expr.(ast.Expression))

	for p.peekTokenIs(token.Comma) {
		p.nextToken()
		p.nextToken()
		if p.curTokenIs(end) {
			return list
		}
		expr := p.parseExpression(priLowest)
		if expr == nil {
			return nil
		}
		list = append(list, expr.(ast.Expression))
	}

	if !p.expectPeek(end) {
		return nil
	}

	return list
}

func (p *Parser) parsePrefixExpression() ast.Expression {
	if p.settings.Debug {
		fmt.Println("parsePrefixExpression")
	}
	expression := &ast.PrefixExpression{
		Token:    p.curToken,
		Operator: p.curToken.Literal,
	}

	p.nextToken()
	expression.Right = p.parseExpression(priPrefix).(ast.Expression)
	return expression
}

func (p *Parser) parseInfixExpression(left ast.Expression) ast.Node {
	if p.settings.Debug {
		fmt.Println("parseInfixExpression")
	}
	expression := &ast.InfixExpression{
		Token:    p.curToken,
		Operator: p.curToken.Literal,
		Left:     left,
	}

	precedence := p.curPrecedence()
	p.nextToken()
	expression.Right = p.parseExpression(precedence).(ast.Expression)
	return expression
}

func (p *Parser) parseGroupedExpression() ast.Expression {
	return p.parseGroupedExpressionC(token.RParen)
}

func (p *Parser) parseGroupedExpressionC(end token.TokenType) ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseGroupedExpressionC")
	}
	p.nextToken()
	exp := p.parseExpression(priLowest).(ast.Expression)

	if !p.expectPeek(end) {
		return nil
	}
	return exp
}

func (p *Parser) parseGroupedExpressionE() ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseGroupedExpressionC")
	}
	p.nextToken()
	return p.parseExpression(priLowest).(ast.Expression)
}
<commit_msg>Catch EOF earlier in block statements<commit_after>package parser

import (
	"fmt"

	"github.com/nitrogen-lang/nitrogen/src/ast"
	"github.com/nitrogen-lang/nitrogen/src/lexer"
	"github.com/nitrogen-lang/nitrogen/src/token"
)

const (
	priLowest      int = iota
	priCompare         // and, or
	priEquals          // ==
	priLessGreater     // > or <
	priSum             // +, -
	priProduct         // *, /
	priPrefix          // -x or !x
	priCall            // myFunction(x)
	priIndex           // array[index]
	priAssign
)

var precedences = map[token.TokenType]int{
	token.LAnd:          priCompare,
	token.LOr:           priCompare,
	token.Equal:         priEquals,
	token.NotEqual:      priEquals,
	token.LessThanEq:    priEquals,
	token.GreaterThanEq: priEquals,
	token.LessThan:      priLessGreater,
	token.GreaterThan:   priLessGreater,
	token.Plus:          priSum,
	token.Dash:          priSum,
	token.BitwiseOr:     priSum,
	token.Carrot:        priSum,
	token.Slash:         priProduct,
	token.Asterisk:      priProduct,
	token.Modulo:        priProduct,
	token.ShiftLeft:     priProduct,
	token.ShiftRight:    priProduct,
	token.BitwiseAnd:    priProduct,
	token.BitwiseAndNot: priProduct,
	token.LParen:        priCall,
	token.LSquare:       priIndex,
	token.Dot:           priIndex,
	token.Assign:        priAssign,
	token.PlusAssign:    priAssign,
	token.MinusAssign:   priAssign,
	token.TimesAssign:   priAssign,
	token.SlashAssign:   priAssign,
	token.ModAssign:     priAssign,
}

type (
	prefixParseFn func() ast.Expression
	infixParseFn  func(ast.Expression) ast.Node
	Settings      struct {
		Debug bool
	}
)

type Parser struct {
	l        *lexer.Lexer
	errors   []string
	settings *Settings

	lastToken token.Token
	curToken  token.Token
	peekToken token.Token

	insertedTokens []token.Token

	prefixParseFns map[token.TokenType]prefixParseFn
	infixParseFns  map[token.TokenType]infixParseFn
}

func New(l *lexer.Lexer, settings *Settings) *Parser {
	if settings == nil {
		settings = &Settings{}
	}

	p := &Parser{
		l:              l,
		settings:       settings,
		errors:         []string{},
		insertedTokens: make([]token.Token, 0, 5),
	}

	// Prefix parsing functions
	p.prefixParseFns = make(map[token.TokenType]prefixParseFn)
	p.registerPrefix(token.Identifier, p.parseIdentifier)
	p.registerPrefix(token.Integer, p.parseIntegerLiteral)
	p.registerPrefix(token.Float, p.parseFloatLiteral)
	p.registerPrefix(token.Nil, p.parseNullLiteral)
	p.registerPrefix(token.String, p.parseStringLiteral)
	p.registerPrefix(token.True, p.parseBoolean)
	p.registerPrefix(token.False, p.parseBoolean)
	p.registerPrefix(token.LSquare, p.parseArrayLiteral)
	p.registerPrefix(token.LBrace, p.parseHashLiteral)
	p.registerPrefix(token.Function, p.parseFunctionLiteral)
	p.registerPrefix(token.Bang, p.parsePrefixExpression)
	p.registerPrefix(token.Dash, p.parsePrefixExpression)
	p.registerPrefix(token.LParen, p.parseGroupedExpression)
	p.registerPrefix(token.If, p.parseIfExpression)
	p.registerPrefix(token.Try, p.parseTryCatch)
	p.registerPrefix(token.Class, p.parseClassLiteral)
	p.registerPrefix(token.New, p.parseMakeExpression)

	// Infix parsing functions
	p.infixParseFns = make(map[token.TokenType]infixParseFn)
	p.registerInfix(token.Plus, p.parseInfixExpression)
	p.registerInfix(token.Dash, p.parseInfixExpression)
	p.registerInfix(token.Slash, p.parseInfixExpression)
	p.registerInfix(token.Asterisk, p.parseInfixExpression)
	p.registerInfix(token.Modulo, p.parseInfixExpression)
	p.registerInfix(token.Equal, p.parseInfixExpression)
	p.registerInfix(token.NotEqual, p.parseInfixExpression)
	p.registerInfix(token.LessThanEq, p.parseInfixExpression)
	p.registerInfix(token.GreaterThanEq, p.parseInfixExpression)
	p.registerInfix(token.LessThan, p.parseInfixExpression)
	p.registerInfix(token.GreaterThan, p.parseInfixExpression)
	p.registerInfix(token.LAnd, p.parseCompareExpression)
	p.registerInfix(token.LOr, p.parseCompareExpression)
	p.registerInfix(token.LParen, p.parseCallExpression)
	p.registerInfix(token.LSquare, p.parseIndexExpression)
	p.registerInfix(token.Dot, p.parseAttributeExpression)
	p.registerInfix(token.Assign, p.parseAssignmentStatement)
	p.registerInfix(token.PlusAssign, p.parseCompoundAssign)
	p.registerInfix(token.MinusAssign, p.parseCompoundAssign)
	p.registerInfix(token.TimesAssign, p.parseCompoundAssign)
	p.registerInfix(token.SlashAssign, p.parseCompoundAssign)
	p.registerInfix(token.ModAssign, p.parseCompoundAssign)
	p.registerInfix(token.ShiftLeft, p.parseInfixExpression)
	p.registerInfix(token.ShiftRight, p.parseInfixExpression)
	p.registerInfix(token.BitwiseAnd, p.parseInfixExpression)
	p.registerInfix(token.BitwiseAndNot, p.parseInfixExpression)
	p.registerInfix(token.BitwiseOr, p.parseInfixExpression)
	p.registerInfix(token.Carrot, p.parseInfixExpression)

	// Read the first two tokens to populate curToken and peekToken
	p.nextToken()
	p.nextToken()

	return p
}

func (p *Parser) Errors() []string {
	return p.errors
}

func (p *Parser) registerPrefix(tt token.TokenType, fn prefixParseFn) {
	p.prefixParseFns[tt] = fn
}

func (p *Parser) registerInfix(tt token.TokenType, fn infixParseFn) {
	p.infixParseFns[tt] = fn
}

func (p *Parser) addError(format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	p.errors = append(p.errors, msg)
}

func (p *Parser) addErrorWithPos(format string, args ...interface{}) {
	if len(args) > 0 {
		args = append([]interface{}{p.curToken.Pos.Line, p.curToken.Pos.Col}, args...)
	} else {
		args = []interface{}{p.curToken.Pos.Line, p.curToken.Pos.Col}
	}
	msg := fmt.Sprintf("at line %d, col %d "+format, args...)
	p.errors = append(p.errors, msg)
}

func (p *Parser) peekError(t token.TokenType) {
	p.addError(
		"at line %d, col %d Incorrect next token. Expected %q, got %q",
		p.peekToken.Pos.Line,
		p.peekToken.Pos.Col,
		t.String(),
		p.peekToken.Type.String(),
	)
}

func (p *Parser) noPrefixParseFnError(t token.TokenType) {
	p.addErrorWithPos("Invalid prefix: %s", t)
}

func (p *Parser) nextToken() {
	p.advanceToken()
	for p.peekTokenIs(token.Comment) {
		p.advancePeekToken()
	}

	if p.curTokenIs(token.Illegal) {
		p.addErrorWithPos("Got illegal token: %s", p.curToken.Literal)
		p.nextToken()
	}

	if p.settings.Debug {
		fmt.Println(p.curToken)
	}
}

func (p *Parser) advanceToken() {
	p.lastToken = p.curToken
	p.curToken = p.peekToken
	p.advancePeekToken()
}

func (p *Parser) advancePeekToken() {
	if len(p.insertedTokens) > 0 {
		p.peekToken = p.insertedTokens[len(p.insertedTokens)-1]
		p.insertedTokens = p.insertedTokens[:len(p.insertedTokens)-1]
	} else {
		p.peekToken = p.l.NextToken()
	}
}

func (p *Parser) insertToken(t token.Token) {
	p.insertedTokens = append(p.insertedTokens, p.peekToken)
	p.peekToken = t
}

func (p *Parser) ParseProgram() *ast.Program {
	if p.settings.Debug {
		fmt.Println("ParseProgram")
	}
	program := &ast.Program{Filename: p.curToken.Filename}
	program.Statements = []ast.Statement{}

	for p.curToken.Type != token.EOF {
		stmt := p.parseStatement()
		if stmt != nil {
			program.Statements = append(program.Statements, stmt)
		}
		if len(p.errors) > 0 {
			return program
		}
		p.nextToken()
	}

	return program
}

func (p *Parser) parseExpression(precedence int) ast.Node {
	if p.settings.Debug {
		fmt.Println("parseExpression")
	}
	prefix := p.prefixParseFns[p.curToken.Type]
	if prefix == nil {
		p.noPrefixParseFnError(p.curToken.Type)
		return nil
	}
	var leftExp ast.Node
	leftExp = prefix()

	for !p.peekTokenIs(token.Semicolon) && precedence < p.peekPrecedence() {
		infix := p.infixParseFns[p.peekToken.Type]
		if infix == nil {
			return leftExp
		}

		p.nextToken()
		leftExp = infix(leftExp.(ast.Expression))
	}

	return leftExp
}

func (p *Parser) parseExpressionStatement() ast.Statement {
	if p.settings.Debug {
		fmt.Println("parseExpressionStatement")
	}
	var stmt ast.Statement
	curToken := p.curToken

	exp := p.parseExpression(priLowest)
	if expStmt, ok := exp.(ast.Statement); ok {
		stmt = expStmt
	} else if exp == nil {
		return nil
	} else {
		stmt = &ast.ExpressionStatement{
			Token:      curToken,
			Expression: exp.(ast.Expression),
		}
	}

	if p.peekTokenIs(token.Semicolon) {
		p.nextToken()
	}

	return stmt
}

func (p *Parser) parseBlockStatements() *ast.BlockStatement {
	if p.settings.Debug {
		fmt.Println("parseBlockStatements")
	}
	block := &ast.BlockStatement{
		Token:      p.curToken,
		Statements: []ast.Statement{},
	}
	p.nextToken()

	for !p.curTokenIs(token.RBrace) && !p.curTokenIs(token.EOF) {
		stmt := p.parseStatement()
		if stmt != nil {
			block.Statements = append(block.Statements, stmt)
		}
		p.nextToken()
	}

	return block
}

func (p *Parser) parseSingleStmtBlock() *ast.BlockStatement {
	if p.settings.Debug {
		fmt.Println("parseSingleStmtBlock")
	}
	block := &ast.BlockStatement{
		Token:      p.curToken,
		Statements: make([]ast.Statement, 1),
	}
	p.nextToken()

	stmt := p.parseStatement()
	if stmt != nil {
		block.Statements[0] = stmt
	}

	return block
}

func (p *Parser) parseExpressionList(end token.TokenType) []ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseExpressionList")
	}
	list := []ast.Expression{}

	if p.peekTokenIs(end) {
		p.nextToken()
		return list
	}

	p.nextToken()
	expr := p.parseExpression(priLowest)
	if expr == nil {
		return nil
	}
	list = append(list, expr.(ast.Expression))

	for p.peekTokenIs(token.Comma) {
		p.nextToken()
		p.nextToken()
		if p.curTokenIs(end) {
			return list
		}
		expr := p.parseExpression(priLowest)
		if expr == nil {
			return nil
		}
		list = append(list, expr.(ast.Expression))
	}

	if !p.expectPeek(end) {
		return nil
	}

	return list
}

func (p *Parser) parsePrefixExpression() ast.Expression {
	if p.settings.Debug {
		fmt.Println("parsePrefixExpression")
	}
	expression := &ast.PrefixExpression{
		Token:    p.curToken,
		Operator: p.curToken.Literal,
	}

	p.nextToken()
	expression.Right = p.parseExpression(priPrefix).(ast.Expression)
	return expression
}

func (p *Parser) parseInfixExpression(left ast.Expression) ast.Node {
	if p.settings.Debug {
		fmt.Println("parseInfixExpression")
	}
	expression := &ast.InfixExpression{
		Token:    p.curToken,
		Operator: p.curToken.Literal,
		Left:     left,
	}

	precedence := p.curPrecedence()
	p.nextToken()
	expression.Right = p.parseExpression(precedence).(ast.Expression)
	return expression
}

func (p *Parser) parseGroupedExpression() ast.Expression {
	return p.parseGroupedExpressionC(token.RParen)
}

func (p *Parser) parseGroupedExpressionC(end token.TokenType) ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseGroupedExpressionC")
	}
	p.nextToken()
	exp := p.parseExpression(priLowest).(ast.Expression)

	if !p.expectPeek(end) {
		return nil
	}
	return exp
}

func (p *Parser) parseGroupedExpressionE() ast.Expression {
	if p.settings.Debug {
		fmt.Println("parseGroupedExpressionC")
	}
	p.nextToken()
	return p.parseExpression(priLowest).(ast.Expression)
}
<|endoftext|>


<commit_before>/*
 *  Copyright (c) 2013 Hualiang Wu <wizawu@gmail.com>
 *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy
 *  of this software and associated documentation files (the "Software"), to
 *  deal in the Software without restriction, including without limitation the
 *  rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 *  sell copies of the Software, and to permit persons to whom the Software is
 *  furnished to do so, subject to the following conditions:
 *
 *  The above copyright notice and this permission notice shall be included in
 *  all copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 *  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 *  IN THE SOFTWARE.
 */

package main

import (
    "flag"
    "strconv"
    "./meepodb"
)

func help() {
    println("PLEASE RUN:\tmeepodb-server [port]")
}

func main() {
    flag.Parse()
    if flag.NArg() != 1 {
        help()
        return
    }
    port, err := strconv.Atoi(flag.Arg(0))
    if err != nil {
        help()
        return
    }

//  meepodb.Reallocate(port)
    meepodb.StartServer(port)
}
<commit_msg>server ok<commit_after>/*
 *  Copyright (c) 2013 Hualiang Wu <wizawu@gmail.com>
 *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy
 *  of this software and associated documentation files (the "Software"), to
 *  deal in the Software without restriction, including without limitation the
 *  rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 *  sell copies of the Software, and to permit persons to whom the Software is
 *  furnished to do so, subject to the following conditions:
 *
 *  The above copyright notice and this permission notice shall be included in
 *  all copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 *  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 *  IN THE SOFTWARE.
 */

package main

import (
    "flag"
    "net"
    "strconv"
    "strings"
    "syscall"
    "./meepodb"
)

func help() {
    println("PLEASE RUN:\tmeepodb-server [port]")
}

func main() {
    /* Get port number from args */
    flag.Parse()
    if flag.NArg() != 1 {
        help()
        return
    }
    port, err := strconv.Atoi(flag.Arg(0))
    if err != nil {
        help()
        return
    }
    /* Check whether the address is defined in config.go */
    addrs, err := net.InterfaceAddrs()
    if err != nil {
        println("Check network interfaces.")
        return
    }
    var ok bool = false
    for _, addr := range addrs {
        s := strings.Split(addr.String(), "/")[0] + ":" + flag.Arg(0)
        for _, svr := range meepodb.SERVERS {
            if s == svr {
                ok = true
                println("Hi, MeepoDB on " + svr)
            }
        }
    }
    if !ok {
        println("Your address is not one of the servers.")
        return
    }

    var dir string = meepodb.DB_DIR
    err = syscall.Chdir(dir)
    if err != nil {
        /* If db does not exist... */
        err = syscall.Mkdir(dir, meepodb.S_IWALL)
        if err != nil {
            panic(err)
        }
    }

//  meepodb.Reallocate(port)
    meepodb.StartServer(port)
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015 Fabrício Godoy
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package memstore

import (
	"strconv"
	"sync"
	"time"

	"gopkg.in/raiqub/data.v0"
	"gopkg.in/raiqub/dot.v1"
)

// A Store provides in-memory key:value cache that expires after defined
// duration of time.
//
// It is a implementation of Store interface.
type Store struct {
	values      map[string]*entry
	lifetime    time.Duration
	isTransient bool
	mutex       sync.RWMutex
	lastgc      time.Time
}

// New creates a new instance of in-memory Store and defines the default
// lifetime for new stored items.
//
// If it is specified to not transient then the stored items lifetime are
// renewed when it is read or written; Otherwise, it is never renewed.
func New(d time.Duration, isTransient bool) *Store {
	return &Store{
		values:      make(map[string]*entry),
		lifetime:    d,
		isTransient: isTransient,
	}
}

// Add adds a new key:value to current store.
//
// Errors:
// DuplicatedKeyError when requested key already exists.
func (s *Store) Add(key string, value interface{}) error {
	switch s.gc() {
	case dot.ReadLocked:
		s.mutex.RUnlock()
		s.mutex.Lock()
	case dot.Unlocked:
		s.mutex.Lock()
	}
	defer s.mutex.Unlock()

	data, err := newEntry(s.lifetime, value)
	if err != nil {
		return err
	}

	if _, ok := s.values[key]; ok {
		return dot.DuplicatedKeyError(key)
	}

	s.values[key] = data
	return nil
}

func (s *Store) atomicInteger(key string, inc int) (int, error) {
	switch s.gc() {
	case dot.ReadLocked:
		s.mutex.RUnlock()
		s.mutex.Lock()
	case dot.Unlocked:
		s.mutex.Lock()
	}
	defer s.mutex.Unlock()

	v, err := s.unsafeGet(key)
	if err != nil {
		data, err := newEntry(s.lifetime, inc)
		if err != nil {
			return 0, err
		}

		s.values[key] = data
		return inc, nil
	}

	var value int
	if err := v.Value(&value); err != nil {
		return 0, err
	}

	value += inc
	v.SetValue(value)

	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}

	return value, nil
}

// Count gets the number of stored values by current instance.
func (s *Store) Count() (int, error) {
	switch s.gc() {
	case dot.WriteLocked:
		s.mutex.Unlock()
	case dot.ReadLocked:
		s.mutex.RUnlock()
	}

	return len(s.values), nil
}

// Decrement atomically gets the value stored by specified key and
// decrements it by one. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) Decrement(key string) (int, error) {
	return s.atomicInteger(key, -1)
}

// DecrementBy atomically gets the value stored by specified key and
// decrements it by value. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) DecrementBy(key string, value int) (int, error) {
	return s.atomicInteger(key, -1*value)
}

// Delete deletes the specified key:value.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Delete(key string) error {
	lckStatus := s.gc()

	_, err := s.unsafeGet(key)
	if err != nil {
		switch lckStatus {
		case dot.WriteLocked:
			s.mutex.Unlock()
		case dot.ReadLocked:
			s.mutex.RUnlock()
		}

		return err
	}

	switch lckStatus {
	case dot.ReadLocked:
		s.mutex.RUnlock()
		s.mutex.Lock()
	case dot.Unlocked:
		s.mutex.Lock()
	}
	defer s.mutex.Unlock()

	delete(s.values, key)
	return nil
}

// Flush deletes any cached value into current instance.
func (s *Store) Flush() error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.values = make(map[string]*entry)
	return nil
}

// Get gets the value stored by specified key.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Get(key string, ref interface{}) error {
	switch s.gc() {
	case dot.WriteLocked:
		s.mutex.Unlock()
		s.mutex.RLock()
	case dot.Unlocked:
		s.mutex.RLock()
	}
	defer s.mutex.RUnlock()

	v, err := s.unsafeGet(key)
	if err != nil {
		return err
	}
	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}

	return v.Value(ref)
}

// GC garbage collects all expired data.
func (s *Store) GC() {
	switch s.gc() {
	case dot.ReadLocked:
		s.mutex.RUnlock()
	case dot.WriteLocked:
		s.mutex.Unlock()
	}
}

func (s *Store) gc() dot.LockStatus {
	// Do not GC intervals smaller than 1/5 of current lifetime
	minInterval := s.lifetime / 5
	if s.lastgc.Add(minInterval).After(time.Now()) {
		return dot.Unlocked
	}

	s.lastgc = time.Now()
	writeLocked := false
	s.mutex.RLock()
	for i := range s.values {
		if s.values[i].IsExpired() {
			if !writeLocked {
				s.mutex.RUnlock()
				s.mutex.Lock()
				writeLocked = true
			}
			// TODO: Investigate how buckets are consolidated
			delete(s.values, i)
		}
	}

	if writeLocked {
		return dot.WriteLocked
	}

	return dot.ReadLocked
}

// Increment atomically gets the value stored by specified key and
// increments it by one. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) Increment(key string) (int, error) {
	return s.atomicInteger(key, 1)
}

// IncrementBy atomically gets the value stored by specified key and
// increments it by value. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) IncrementBy(key string, value int) (int, error) {
	return s.atomicInteger(key, value)
}

// Set sets the value of specified key.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Set(key string, value interface{}) error {
	switch s.gc() {
	case dot.ReadLocked:
		s.mutex.RUnlock()
		s.mutex.Lock()
	case dot.Unlocked:
		s.mutex.Lock()
	}
	defer s.mutex.Unlock()

	v, err := s.unsafeGet(key)
	if err != nil {
		return err
	}

	v.SetValue(value)

	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}
	return nil
}

// SetLifetime modifies the lifetime for new stored items or for existing items
// when it is read or written.
//
// Errors:
// NotSupportedError when ScopeNew is specified.
func (s *Store) SetLifetime(d time.Duration, scope data.LifetimeScope) error {
	switch scope {
	case data.ScopeAll:
		switch s.gc() {
		case dot.ReadLocked:
			s.mutex.RUnlock()
			s.mutex.Lock()
		case dot.Unlocked:
			s.mutex.Lock()
		}
		defer s.mutex.Unlock()

		for _, v := range s.values {
			v.SetLifetime(d)
		}
	case data.ScopeNewAndUpdated:
		s.mutex.RLock()
		defer s.mutex.RUnlock()
	case data.ScopeNew:
		return dot.NotSupportedError("ScopeNew")
	default:
		return dot.NotSupportedError(strconv.Itoa(int(scope)))
	}

	s.lifetime = d
	return nil
}

// SetTransient defines whether should extends expiration of stored value when
// it is read or written.
func (s *Store) SetTransient(value bool) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	s.isTransient = value
}

// unsafeGet gets one entry instance from its key without locking.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) unsafeGet(key string) (*entry, error) {
	v, ok := s.values[key]
	if !ok {
		return nil, dot.InvalidKeyError(key)
	}
	return v, nil
}

var _ data.Store = (*Store)(nil)
<commit_msg>Improvement to run GC in parallel (memstore)<commit_after>/*
 * Copyright 2015 Fabrício Godoy
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package memstore

import (
	"strconv"
	"sync"
	"time"

	"gopkg.in/raiqub/data.v0"
	"gopkg.in/raiqub/dot.v1"
)

// A Store provides in-memory key:value cache that expires after defined
// duration of time.
//
// It is a implementation of Store interface.
type Store struct {
	values      map[string]*entry
	lifetime    time.Duration
	isTransient bool
	mutex       sync.RWMutex
	gcRunning   bool
}

// New creates a new instance of in-memory Store and defines the default
// lifetime for new stored items.
//
// If it is specified to not transient then the stored items lifetime are
// renewed when it is read or written; Otherwise, it is never renewed.
func New(d time.Duration, isTransient bool) *Store {
	return &Store{
		values:      make(map[string]*entry),
		lifetime:    d,
		isTransient: isTransient,
	}
}

// Add adds a new key:value to current store.
//
// Errors:
// DuplicatedKeyError when requested key already exists.
func (s *Store) Add(key string, value interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	data, err := newEntry(s.lifetime, value)
	if err != nil {
		return err
	}

	if _, ok := s.values[key]; ok {
		return dot.DuplicatedKeyError(key)
	}

	if !s.gcRunning {
		go s.gc()
	}
	s.values[key] = data
	return nil
}

func (s *Store) atomicInteger(key string, inc int) (int, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	v, err := s.unsafeGet(key)
	if err != nil {
		data, err := newEntry(s.lifetime, inc)
		if err != nil {
			return 0, err
		}

		if !s.gcRunning {
			go s.gc()
		}
		s.values[key] = data
		return inc, nil
	}

	var value int
	if err := v.Value(&value); err != nil {
		return 0, err
	}

	value += inc
	v.SetValue(value)

	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}

	return value, nil
}

// Count gets the number of stored values by current instance.
func (s *Store) Count() (int, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	return len(s.values), nil
}

// Decrement atomically gets the value stored by specified key and
// decrements it by one. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) Decrement(key string) (int, error) {
	return s.atomicInteger(key, -1)
}

// DecrementBy atomically gets the value stored by specified key and
// decrements it by value. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) DecrementBy(key string, value int) (int, error) {
	return s.atomicInteger(key, -1*value)
}

// Delete deletes the specified key:value.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Delete(key string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	_, err := s.unsafeGet(key)
	if err != nil {
		return err
	}

	delete(s.values, key)
	return nil
}

// Flush deletes any cached value into current instance.
func (s *Store) Flush() error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	s.values = make(map[string]*entry)
	return nil
}

// Get gets the value stored by specified key.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Get(key string, ref interface{}) error {
	if s.isTransient {
		s.mutex.RLock()
		defer s.mutex.RUnlock()
	} else {
		s.mutex.Lock()
		defer s.mutex.Unlock()
	}

	v, err := s.unsafeGet(key)
	if err != nil {
		return err
	}
	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}

	return v.Value(ref)
}

// GC garbage collects all expired data.
func (s *Store) GC() {
}

func (s *Store) gc() {
	s.mutex.Lock()
	if s.gcRunning {
		s.mutex.Unlock()
		return
	}

	// Schedule GC at 1/5 intervals of current lifetime.
	interval := s.lifetime / 5
	s.gcRunning = true
	s.mutex.Unlock()

	for {
		<-time.After(interval)

		writeLocked := false
		s.mutex.RLock()
		for i := range s.values {
			if s.values[i].IsExpired() {
				if !writeLocked {
					s.mutex.RUnlock()
					s.mutex.Lock()
					writeLocked = true
				}
				// TODO: Investigate how buckets are consolidated
				delete(s.values, i)
			}
		}

		interval = s.lifetime / 5
		isEmpty := len(s.values) == 0
		if isEmpty {
			s.gcRunning = false
		}
		if writeLocked {
			s.mutex.Unlock()
		} else {
			s.mutex.RUnlock()
		}

		if isEmpty {
			return
		}
	}
}

// Increment atomically gets the value stored by specified key and
// increments it by one. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) Increment(key string) (int, error) {
	return s.atomicInteger(key, 1)
}

// IncrementBy atomically gets the value stored by specified key and
// increments it by value. If the key does not exist, it is created.
//
// Errors:
// InvalidTypeError when the value stored at key is not integer.
func (s *Store) IncrementBy(key string, value int) (int, error) {
	return s.atomicInteger(key, value)
}

// Set sets the value of specified key.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) Set(key string, value interface{}) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	v, err := s.unsafeGet(key)
	if err != nil {
		return err
	}

	v.SetValue(value)

	if !s.isTransient {
		v.SetLifetime(s.lifetime)
		v.Hit()
	}
	return nil
}

// SetLifetime modifies the lifetime for new stored items or for existing items
// when it is read or written.
//
// Errors:
// NotSupportedError when ScopeNew is specified.
func (s *Store) SetLifetime(d time.Duration, scope data.LifetimeScope) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	switch scope {
	case data.ScopeAll:
		for _, v := range s.values {
			v.SetLifetime(d)
		}
	case data.ScopeNewAndUpdated:
	case data.ScopeNew:
		return dot.NotSupportedError("ScopeNew")
	default:
		return dot.NotSupportedError(strconv.Itoa(int(scope)))
	}

	s.lifetime = d
	return nil
}

// SetTransient defines whether should extends expiration of stored value when
// it is read or written.
func (s *Store) SetTransient(value bool) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	s.isTransient = value
}

// unsafeGet gets one entry instance from its key without locking.
//
// Errors:
// InvalidKeyError when requested key could not be found.
func (s *Store) unsafeGet(key string) (*entry, error) {
	v, ok := s.values[key]
	if !ok {
		return nil, dot.InvalidKeyError(key)
	}
	return v, nil
}

var _ data.Store = (*Store)(nil)
<|endoftext|>


<commit_before>package mesos

import (
	"fmt"
	"log"

	"github.com/CiscoCloud/mesos-consul/registry"
)

type registryState struct {
	services	map[string]bool
	followers	map[string]bool
	masters		map[string]MesosHost
}

var currentState = registryState{
	services:	make(map[string]bool),
	followers:	make(map[string]bool),
	masters:	make(map[string]MesosHost),
}

func (m *Mesos) RegisterHosts(sj StateJSON) {
	log.Print("Running RegisterHosts")

	// Register followers
	for _, f := range sj.Followers {
		id := fmt.Sprintf("%s:%s", f.Id, f.Hostname)
		if _, ok := currentState.followers[id]; ok {
			log.Printf("Id found: '%s': Not registering", id)
			currentState.followers[id] = true
		} else {
			currentState.followers[id] = true

			s := new(registry.Service)

			host, port := parsePID(f.Pid)

			s.ID	= id
			s.Name	= "mesos"
			s.Port	= toPort(port)
			s.IP	= toIP(host)
			s.Tags	= []string{
				"follower",
				}

			log.Print("Registering: ", id)

			err := m.registry.Register(s)
			if err != nil {
				log.Print(err)
			}
		}
	}

	// Register masters
	mas := m.getMasters()
	for _, ma := range mas {
		id := fmt.Sprintf("%s:%s", ma.host, ma.port)
		if master, ok := currentState.masters[id]; ok {
			// Master has been found. Only update if the leader status has changed
			if ma.isLeader ==  master.isLeader {
				master.isRegistered = true
				log.Printf("Master state unchanged, not registering: '%s'", id)
				continue
			}
		}

		ma.isRegistered = true
		currentState.masters[id] = ma

		s := new(registry.Service)
		s.ID	= fmt.Sprintf("mesos:%s:%s", ma.host, ma.port)
		s.Name	= "mesos"
		s.Port	= toPort(ma.port)
		s.IP	= toIP(ma.host)
		if ma.isLeader {
			s.Tags	= []string{
				"master",
				}
		} else {
			s.Tags	= []string{
				"master",
				"leader",
				}
		}

		log.Print("Registering: ", s.ID)
		err := m.registry.Register(s)
		if err != nil {
			log.Print(err)
		}
	}
}

func (m *Mesos) register(s *registry.Service) {
	currentState.services[s.ID] = true

	log.Print("Registering ", s.ID)

	err := m.registry.Register(s)
	if err != nil {
		log.Print(err)
	}
}

// deregister items that have gone away
//
func (m *Mesos) deregister() {
	// Services
	for s, b := range currentState.services {
		if !b {
			log.Print("Deregistering ", s)
			m.registry.Deregister(&registry.Service{
				ID:		s,
				})

			delete(currentState.services, s)
		} else {
			currentState.services[s] = false
		}
	}

	// Followers
	for id, isRegistered := range currentState.followers {
		if !isRegistered {
			log.Print("Deregistering ", id)
			m.registry.Deregister(&registry.Service{
				ID:		id,
				})
			delete(currentState.followers, id)
		} else {
			currentState.followers[id] = false
		}
	}

	// Masters
	for id, master := range currentState.masters {
		if !master.isRegistered {
			log.Print("Deregistering ", id)
			m.registry.Deregister(&registry.Service{
				ID:		id,
				})
			delete(currentState.masters, id)
		} else {
			master.isRegistered = false
		}
	}
}
<commit_msg>Check if service is already registered<commit_after>package mesos

import (
	"fmt"
	"log"

	"github.com/CiscoCloud/mesos-consul/registry"
)

type registryState struct {
	services	map[string]bool
	followers	map[string]bool
	masters		map[string]MesosHost
}

var currentState = registryState{
	services:	make(map[string]bool),
	followers:	make(map[string]bool),
	masters:	make(map[string]MesosHost),
}

func (m *Mesos) RegisterHosts(sj StateJSON) {
	log.Print("Running RegisterHosts")

	// Register followers
	for _, f := range sj.Followers {
		id := fmt.Sprintf("%s:%s", f.Id, f.Hostname)
		if _, ok := currentState.followers[id]; ok {
			log.Printf("Id found: '%s': Not registering", id)
			currentState.followers[id] = true
		} else {
			currentState.followers[id] = true

			s := new(registry.Service)

			host, port := parsePID(f.Pid)

			s.ID	= id
			s.Name	= "mesos"
			s.Port	= toPort(port)
			s.IP	= toIP(host)
			s.Tags	= []string{
				"follower",
				}

			log.Print("Registering: ", id)

			err := m.registry.Register(s)
			if err != nil {
				log.Print(err)
			}
		}
	}

	// Register masters
	mas := m.getMasters()
	for _, ma := range mas {
		id := fmt.Sprintf("%s:%s", ma.host, ma.port)
		if master, ok := currentState.masters[id]; ok {
			// Master has been found. Only update if the leader status has changed
			if ma.isLeader ==  master.isLeader {
				master.isRegistered = true
				log.Printf("Master state unchanged, not registering: '%s'", id)
				continue
			}
		}

		ma.isRegistered = true
		currentState.masters[id] = ma

		s := new(registry.Service)
		s.ID	= fmt.Sprintf("mesos:%s:%s", ma.host, ma.port)
		s.Name	= "mesos"
		s.Port	= toPort(ma.port)
		s.IP	= toIP(ma.host)
		if ma.isLeader {
			s.Tags	= []string{
				"master",
				}
		} else {
			s.Tags	= []string{
				"master",
				"leader",
				}
		}

		log.Print("Registering: ", s.ID)
		err := m.registry.Register(s)
		if err != nil {
			log.Print(err)
		}
	}
}

func (m *Mesos) register(s *registry.Service) {
	if currentState.services[s.ID] {
		log.Printf("Service found. Not registering: %s", s.ID)
		currentState.services[s.ID] = true
		return
	}

	log.Print("Registering ", s.ID)

	currentState.services[s.ID] = true
	err := m.registry.Register(s)
	if err != nil {
		log.Print(err)
	}
}

// deregister items that have gone away
//
func (m *Mesos) deregister() {
	// Services
	for s, b := range currentState.services {
		if !b {
			log.Print("Deregistering ", s)
			m.registry.Deregister(&registry.Service{
				ID:		s,
				})

			delete(currentState.services, s)
		} else {
			currentState.services[s] = false
		}
	}

	// Followers
	for id, isRegistered := range currentState.followers {
		if !isRegistered {
			log.Print("Deregistering ", id)
			m.registry.Deregister(&registry.Service{
				ID:		id,
				})
			delete(currentState.followers, id)
		} else {
			currentState.followers[id] = false
		}
	}

	// Masters
	for id, master := range currentState.masters {
		if !master.isRegistered {
			log.Print("Deregistering ", id)
			m.registry.Deregister(&registry.Service{
				ID:		id,
				})
			delete(currentState.masters, id)
		} else {
			master.isRegistered = false
		}
	}
}
<|endoftext|>


<commit_before>package containers

import (
	"atlantis/crypto"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/jigish/go-dockerclient"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
)

var (
	dockerIdRegexp = regexp.MustCompile("^[A-Za-z0-9]+$")
	dockerLock     = sync.Mutex{}
	dockerClient   *docker.Client
)

type Container types.Container
type SSHCmd []string

func DockerInit() (err error) {
	dockerClient, err = docker.NewClient("unix:///var/run/docker.sock")
	return
}

func pretending() bool {
	return os.Getenv("SUPERVISOR_PRETEND") != ""
}

func (s SSHCmd) Execute() error {
	if pretending() {
		log.Printf("[pretend] ssh %s", strings.Join(s, " "))
		return nil
	}
	log.Printf("ssh %s", strings.Join(s, " "))
	cmd := exec.Command("ssh", s...)
	output, err := cmd.CombinedOutput()
	log.Printf("-> %s", output)
	if err != nil {
		log.Println("-> Error:", err)
	}
	return err
}

func removeExited() {
	if pretending() {
		return
	}
	dockerLock.Lock()
	defer dockerLock.Unlock()
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Printf("[RemoveExited] could not list containers: %v", err)
		return
	}
	for _, cont := range containers {
		log.Printf("[RemoveExited] checking %s (%v) : %s", cont.ID, cont.Names, cont.Status)
		if !strings.HasPrefix(cont.Status, "Exit") {
			continue
		}
		log.Printf("[RemoveExited] remove %s (%v)", cont.ID, cont.Names)
		err := dockerClient.RemoveContainer(cont.ID)
		if err != nil {
			log.Printf("[RemoveExited] -> error: %v", err)
		} else {
			log.Printf("[RemoveExited] -> success", err)
		}
	}
}

func restartGhost() {
	if pretending() {
		return
	}
	dockerLock.Lock()
	defer dockerLock.Unlock()
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Printf("[RestartGhost] could not list containers: %v", err)
		return
	}
	for _, cont := range containers {
		log.Printf("[RestartGhost] checking %s (%v) : %s", cont.ID, cont.Names, cont.Status)
		if !strings.HasPrefix(cont.Status, "Ghost") {
			continue
		}
		log.Printf("[RestartGhost] restart %s (%v)", cont.ID, cont.Names)
		err := dockerClient.RestartContainer(cont.ID, 0)
		if err != nil {
			log.Printf("[RestartGhost] -> error: %v", err)
		} else {
			log.Printf("[RestartGhost] -> success", err)
		}
	}
}

func (c *Container) dockerCfgs(repo string) (*docker.Config, *docker.HostConfig) {
	// get env cfg
	envs := []string{
		"ATLANTIS=true",
		fmt.Sprintf("CONTAINER_ID=%s", c.Id),
		fmt.Sprintf("CONTAINER_HOST=%s", c.Host),
		fmt.Sprintf("CONTAINER_ENV=%s", c.Env),
		fmt.Sprintf("HTTP_PORT=%d", c.PrimaryPort),
		fmt.Sprintf("SSHD_PORT=%d", c.SSHPort),
	}
	if c.Manifest.Deps != nil {
		for name, value := range c.Manifest.Deps {
			envs = append(envs, fmt.Sprintf("%s=%s", name, crypto.Decrypt([]byte(value))))
		}
	}

	// get port cfg
	exposedPorts := map[docker.Port]struct{}{}
	portBindings := map[docker.Port][]docker.PortBinding{}
	sPrimaryPort := fmt.Sprintf("%d", c.PrimaryPort)
	dPrimaryPort := docker.NewPort("tcp", sPrimaryPort)
	exposedPorts[dPrimaryPort] = struct{}{}
	portBindings[dPrimaryPort] = []docker.PortBinding{docker.PortBinding{
		HostIp:   "",
		HostPort: fmt.Sprintf("%d", sPrimaryPort),
	}}
	sSSHPort := fmt.Sprintf("%d", c.SSHPort)
	dSSHPort := docker.NewPort("tcp", sSSHPort)
	exposedPorts[dSSHPort] = struct{}{}
	portBindings[dSSHPort] = []docker.PortBinding{docker.PortBinding{
		HostIp:   "",
		HostPort: sSSHPort,
	}}
	for i, port := range c.SecondaryPorts {
		sPort := fmt.Sprintf("%d", port)
		dPort := docker.NewPort("tcp", sPort)
		exposedPorts[dPort] = struct{}{}
		portBindings[dPort] = []docker.PortBinding{docker.PortBinding{
			HostIp:   "",
			HostPort: sPort,
		}}
		envs = append(envs, fmt.Sprintf("SECONDARY_PORT%d=%d", i, port))
	}

	// setup actual cfg
	dCfg := &docker.Config{
		CpuShares:    int64(c.Manifest.CPUShares),
		Memory:       int64(c.Manifest.MemoryLimit) * int64(1024*1024), // this is in bytes
		MemorySwap:   int64(-1),                                        // -1 turns swap off
		ExposedPorts: exposedPorts,
		Env:          envs,
		Cmd:          []string{}, // images already specify run command
		Image:        fmt.Sprintf("%s/%s", RegistryHost, repo),
		Volumes: map[string]struct{}{
			fmt.Sprintf("/var/log/atlantis/containers/%s:/var/log/atlantis/syslog", c.Id): struct{}{},
		},
	}
	dHostCfg := &docker.HostConfig{
		PortBindings: portBindings,
		LxcConf:      []docker.KeyValuePair{},
	}
	return dCfg, dHostCfg
}

// Deploy the given app+sha with the dependencies defined in deps. This will spin up a new docker container.
func (c *Container) Deploy(host, app, sha, env string) error {
	c.Host = host
	c.App = app
	c.Sha = sha
	c.Env = env
	dRepo := fmt.Sprintf("apps/%s-%s", c.App, c.Sha)
	if pretending() {
	} else {
	}
	// Pull docker container
	if pretending() {
		log.Printf("[pretend] deploy %s with %s @ %s...", c.Id, c.App, c.Sha)
		log.Printf("[pretend] docker pull %s/%s", RegistryHost, dRepo)
		log.Printf("[pretend] docker run %s/%s", RegistryHost, dRepo)
		c.DockerId = fmt.Sprintf("pretend-docker-id-%d", c.PrimaryPort)
	} else {
		log.Printf("deploy %s with %s @ %s...", c.Id, c.App, c.Sha)
		dockerLock.Lock()
		err := dockerClient.PullImage(docker.PullImageOptions{Repository: dRepo, Registry: RegistryHost},
			os.Stdout)
		dockerLock.Unlock()
		if err != nil {
			return err
		}

		err = os.MkdirAll(fmt.Sprintf("/var/log/atlantis/containers/%s", c.Id), 0755)
		if err != nil {
			return err
		}

		// create docker container
		dCfg, dHostCfg := c.dockerCfgs(dRepo)
		dockerLock.Lock()
		dCont, err := dockerClient.CreateContainer(docker.CreateContainerOptions{Name: c.Id}, dCfg)
		dockerLock.Unlock()
		if err != nil {
			return err
		}
		c.DockerId = dCont.ID

		// start docker container
		dockerLock.Lock()
		err = dockerClient.StartContainer(c.DockerId, dHostCfg)
		dockerLock.Unlock()
		if err != nil {
			return err
		}
	}
	save() // save here because this is when we know the deployed container is actually alive
	return nil
}

// Teardown the container. This will kill the docker container but will not free the ports/containers
func (c *Container) teardown() {
	if pretending() {
		log.Printf("[pretend] teardown %s...", c.Id)
		return
	} else {
		log.Printf("teardown %s...", c.Id)
	}
	defer removeExited()
	dockerLock.Lock()
	err := dockerClient.KillContainer(c.DockerId)
	dockerLock.Unlock()
	if err != nil {
		log.Printf("failed to teardown %s: %v", c.Id, err)
	}
}

// This calls the Teardown(id string) method to ensure that the ports/containers are freed. That will in turn
// call c.teardown(id string)
func (c *Container) Teardown() {
	Teardown(c.Id)
}

func (c *Container) AuthorizeSSHUser(user, publicKey string) error {
	// copy file to container
	// rebuild authorize_keys
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		fmt.Sprintf("echo \"%s\" >/root/.ssh/authorized_keys.d/%s.pub && rebuild_authorized_keys", publicKey,
			user)}.Execute()
}

func (c *Container) DeauthorizeSSHUser(user string) error {
	// delete file from container
	// rebuild authorize_keys
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		fmt.Sprintf("rm /root/.ssh/authorized_keys.d/%s.pub && rebuild_authorized_keys",
			user)}.Execute()
}

func (c *Container) SetMaintenance(maint bool) error {
	if maint {
		// touch /etc/maint
		return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
			"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
			"touch /etc/maint"}.Execute()
	}
	// rm -f /etc/maint
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		"rm -f /etc/maint"}.Execute()
}
<commit_msg>remove unused param<commit_after>package containers

import (
	"atlantis/crypto"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/jigish/go-dockerclient"
	"log"
	"os"
	"os/exec"
	"regexp"
	"strings"
	"sync"
)

var (
	dockerIdRegexp = regexp.MustCompile("^[A-Za-z0-9]+$")
	dockerLock     = sync.Mutex{}
	dockerClient   *docker.Client
)

type Container types.Container
type SSHCmd []string

func DockerInit() (err error) {
	dockerClient, err = docker.NewClient("unix:///var/run/docker.sock")
	return
}

func pretending() bool {
	return os.Getenv("SUPERVISOR_PRETEND") != ""
}

func (s SSHCmd) Execute() error {
	if pretending() {
		log.Printf("[pretend] ssh %s", strings.Join(s, " "))
		return nil
	}
	log.Printf("ssh %s", strings.Join(s, " "))
	cmd := exec.Command("ssh", s...)
	output, err := cmd.CombinedOutput()
	log.Printf("-> %s", output)
	if err != nil {
		log.Println("-> Error:", err)
	}
	return err
}

func removeExited() {
	if pretending() {
		return
	}
	dockerLock.Lock()
	defer dockerLock.Unlock()
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Printf("[RemoveExited] could not list containers: %v", err)
		return
	}
	for _, cont := range containers {
		log.Printf("[RemoveExited] checking %s (%v) : %s", cont.ID, cont.Names, cont.Status)
		if !strings.HasPrefix(cont.Status, "Exit") {
			continue
		}
		log.Printf("[RemoveExited] remove %s (%v)", cont.ID, cont.Names)
		err := dockerClient.RemoveContainer(cont.ID)
		if err != nil {
			log.Printf("[RemoveExited] -> error: %v", err)
		} else {
			log.Printf("[RemoveExited] -> success")
		}
	}
}

func restartGhost() {
	if pretending() {
		return
	}
	dockerLock.Lock()
	defer dockerLock.Unlock()
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{All: true})
	if err != nil {
		log.Printf("[RestartGhost] could not list containers: %v", err)
		return
	}
	for _, cont := range containers {
		log.Printf("[RestartGhost] checking %s (%v) : %s", cont.ID, cont.Names, cont.Status)
		if !strings.HasPrefix(cont.Status, "Ghost") {
			continue
		}
		log.Printf("[RestartGhost] restart %s (%v)", cont.ID, cont.Names)
		err := dockerClient.RestartContainer(cont.ID, 0)
		if err != nil {
			log.Printf("[RestartGhost] -> error: %v", err)
		} else {
			log.Printf("[RestartGhost] -> success")
		}
	}
}

func (c *Container) dockerCfgs(repo string) (*docker.Config, *docker.HostConfig) {
	// get env cfg
	envs := []string{
		"ATLANTIS=true",
		fmt.Sprintf("CONTAINER_ID=%s", c.Id),
		fmt.Sprintf("CONTAINER_HOST=%s", c.Host),
		fmt.Sprintf("CONTAINER_ENV=%s", c.Env),
		fmt.Sprintf("HTTP_PORT=%d", c.PrimaryPort),
		fmt.Sprintf("SSHD_PORT=%d", c.SSHPort),
	}
	if c.Manifest.Deps != nil {
		for name, value := range c.Manifest.Deps {
			envs = append(envs, fmt.Sprintf("%s=%s", name, crypto.Decrypt([]byte(value))))
		}
	}

	// get port cfg
	exposedPorts := map[docker.Port]struct{}{}
	portBindings := map[docker.Port][]docker.PortBinding{}
	sPrimaryPort := fmt.Sprintf("%d", c.PrimaryPort)
	dPrimaryPort := docker.NewPort("tcp", sPrimaryPort)
	exposedPorts[dPrimaryPort] = struct{}{}
	portBindings[dPrimaryPort] = []docker.PortBinding{docker.PortBinding{
		HostIp:   "",
		HostPort: fmt.Sprintf("%d", sPrimaryPort),
	}}
	sSSHPort := fmt.Sprintf("%d", c.SSHPort)
	dSSHPort := docker.NewPort("tcp", sSSHPort)
	exposedPorts[dSSHPort] = struct{}{}
	portBindings[dSSHPort] = []docker.PortBinding{docker.PortBinding{
		HostIp:   "",
		HostPort: sSSHPort,
	}}
	for i, port := range c.SecondaryPorts {
		sPort := fmt.Sprintf("%d", port)
		dPort := docker.NewPort("tcp", sPort)
		exposedPorts[dPort] = struct{}{}
		portBindings[dPort] = []docker.PortBinding{docker.PortBinding{
			HostIp:   "",
			HostPort: sPort,
		}}
		envs = append(envs, fmt.Sprintf("SECONDARY_PORT%d=%d", i, port))
	}

	// setup actual cfg
	dCfg := &docker.Config{
		CpuShares:    int64(c.Manifest.CPUShares),
		Memory:       int64(c.Manifest.MemoryLimit) * int64(1024*1024), // this is in bytes
		MemorySwap:   int64(-1),                                        // -1 turns swap off
		ExposedPorts: exposedPorts,
		Env:          envs,
		Cmd:          []string{}, // images already specify run command
		Image:        fmt.Sprintf("%s/%s", RegistryHost, repo),
		Volumes: map[string]struct{}{
			fmt.Sprintf("/var/log/atlantis/containers/%s:/var/log/atlantis/syslog", c.Id): struct{}{},
		},
	}
	dHostCfg := &docker.HostConfig{
		PortBindings: portBindings,
		LxcConf:      []docker.KeyValuePair{},
	}
	return dCfg, dHostCfg
}

// Deploy the given app+sha with the dependencies defined in deps. This will spin up a new docker container.
func (c *Container) Deploy(host, app, sha, env string) error {
	c.Host = host
	c.App = app
	c.Sha = sha
	c.Env = env
	dRepo := fmt.Sprintf("apps/%s-%s", c.App, c.Sha)
	if pretending() {
	} else {
	}
	// Pull docker container
	if pretending() {
		log.Printf("[pretend] deploy %s with %s @ %s...", c.Id, c.App, c.Sha)
		log.Printf("[pretend] docker pull %s/%s", RegistryHost, dRepo)
		log.Printf("[pretend] docker run %s/%s", RegistryHost, dRepo)
		c.DockerId = fmt.Sprintf("pretend-docker-id-%d", c.PrimaryPort)
	} else {
		log.Printf("deploy %s with %s @ %s...", c.Id, c.App, c.Sha)
		dockerLock.Lock()
		err := dockerClient.PullImage(docker.PullImageOptions{Repository: dRepo, Registry: RegistryHost},
			os.Stdout)
		dockerLock.Unlock()
		if err != nil {
			return err
		}

		err = os.MkdirAll(fmt.Sprintf("/var/log/atlantis/containers/%s", c.Id), 0755)
		if err != nil {
			return err
		}

		// create docker container
		dCfg, dHostCfg := c.dockerCfgs(dRepo)
		dockerLock.Lock()
		dCont, err := dockerClient.CreateContainer(docker.CreateContainerOptions{Name: c.Id}, dCfg)
		dockerLock.Unlock()
		if err != nil {
			return err
		}
		c.DockerId = dCont.ID

		// start docker container
		dockerLock.Lock()
		err = dockerClient.StartContainer(c.DockerId, dHostCfg)
		dockerLock.Unlock()
		if err != nil {
			return err
		}
	}
	save() // save here because this is when we know the deployed container is actually alive
	return nil
}

// Teardown the container. This will kill the docker container but will not free the ports/containers
func (c *Container) teardown() {
	if pretending() {
		log.Printf("[pretend] teardown %s...", c.Id)
		return
	} else {
		log.Printf("teardown %s...", c.Id)
	}
	defer removeExited()
	dockerLock.Lock()
	err := dockerClient.KillContainer(c.DockerId)
	dockerLock.Unlock()
	if err != nil {
		log.Printf("failed to teardown %s: %v", c.Id, err)
	}
}

// This calls the Teardown(id string) method to ensure that the ports/containers are freed. That will in turn
// call c.teardown(id string)
func (c *Container) Teardown() {
	Teardown(c.Id)
}

func (c *Container) AuthorizeSSHUser(user, publicKey string) error {
	// copy file to container
	// rebuild authorize_keys
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		fmt.Sprintf("echo \"%s\" >/root/.ssh/authorized_keys.d/%s.pub && rebuild_authorized_keys", publicKey,
			user)}.Execute()
}

func (c *Container) DeauthorizeSSHUser(user string) error {
	// delete file from container
	// rebuild authorize_keys
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		fmt.Sprintf("rm /root/.ssh/authorized_keys.d/%s.pub && rebuild_authorized_keys",
			user)}.Execute()
}

func (c *Container) SetMaintenance(maint bool) error {
	if maint {
		// touch /etc/maint
		return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
			"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
			"touch /etc/maint"}.Execute()
	}
	// rm -f /etc/maint
	return SSHCmd{"-p", fmt.Sprintf("%d", c.SSHPort), "-i", "/opt/atlantis/supervisor/master_id_rsa", "-o",
		"UserKnownHostsFile=/dev/null", "-o", "StrictHostKeyChecking=no", "root@localhost",
		"rm -f /etc/maint"}.Execute()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"container/list"
	"database/sql"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"html"
	"html/template"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"
)

type Import struct {
	Id         uint64 `col:"id"`
	PostId     uint64 `col:"post"`
	Source     string `col:"source"`
	Identifier string `col:"identifier"`
}

func NewImport() *Import {
	return &Import{0, 0, "", ""}
}

func ImportBySourceIdentifier(source, identifier string) (*Import, error) {
	row := db.QueryRow("SELECT id, post, source, identifier FROM import WHERE source = $1 AND identifier = $2 LIMIT 1",
		source, identifier)

	var id, postid uint64
	err := row.Scan(&id, &postid, &source, &identifier)
	if err != nil {
		return nil, err
	}

	i := &Import{id, postid, source, identifier}
	return i, nil
}

func (im *Import) Save() error {
	return db.Save(im, "import")
}

type Mutation struct {
	Start int
	End   int
	Html  string
}

type MutationList []Mutation

func (ml MutationList) Len() int {
	return len(ml)
}

func (ml MutationList) Less(i, j int) bool {
	mutI := ml[i]
	mutJ := ml[j]
	if mutI.Start < mutJ.Start {
		return true
	}
	if mutI.Start == mutJ.Start && mutI.End < mutJ.End {
		return true
	}
	return false
}

func (ml MutationList) Swap(i, j int) {
	ml[i], ml[j] = ml[j], ml[i]
}

func indicesForEntity(ent map[string]interface{}) (int, int) {
	indices := ent["indices"].([]interface{})
	s, e := indices[0].(float64), indices[1].(float64)
	return int(s), int(e)
}

func makeTweetMutations(data map[string]interface{}) MutationList {
	text := data["text"].(string)
	ents := data["entities"].(map[string]interface{})

	mutations := list.New()
	for _, entIf := range ents["user_mentions"].([]interface{}) {
		ent := entIf.(map[string]interface{})
		screenName := html.EscapeString(ent["screen_name"].(string))
		html := fmt.Sprintf(`<a href="https://twitter.com/%s" title="%s">@%s</a>`,
			screenName, html.EscapeString(ent["name"].(string)), screenName)
		start, end := indicesForEntity(ent)
		mutations.PushBack(Mutation{start, end, html})
	}
	for _, entIf := range ents["hashtags"].([]interface{}) {
		ent := entIf.(map[string]interface{})
		tagText := ent["text"].(string)
		html := fmt.Sprintf(`<a href="https://twitter.com/search?q=%%23%s">#%s</a>`,
			tagText, tagText)
		start, end := indicesForEntity(ent)
		mutations.PushBack(Mutation{start, end, html})
	}
	for _, urlEnts := range []interface{}{ents["urls"], ents["media"]} {
		if urlEnts == nil {
			continue
		}
		for _, entIf := range urlEnts.([]interface{}) {
			ent := entIf.(map[string]interface{})
			url := ent["expanded_url"].(string)
			if url == "" {
				url = ent["url"].(string)
			}
			linkText := ent["display_url"].(string)
			html := fmt.Sprintf(`<a href="%s">%s</a>`, url, linkText)
			start, end := indicesForEntity(ent)
			mutations.PushBack(Mutation{start, end, html})
		}
	}

	// We don't strictly need to regexp this of course but the strings package
	// won't find *all* instances of a substring, only the first or last.
	// Sadface that we can't just use lookahead assertion too.
	ampsRE, _ := regexp.Compile(`&`)
	amps := ampsRE.FindAllStringIndex(text, -1)
	for _, ampIndices := range amps {
		rest := text[ampIndices[1]:]
		matched, _ := regexp.MatchString("^(?:lt|gt|amp);", rest)
		if !matched {
			mutations.PushBack(Mutation{ampIndices[0], ampIndices[1], "&"})
		}
	}

	nlRE, _ := regexp.Compile(`\n`)
	nls := nlRE.FindAllStringIndex(text, -1)
	for _, nlIndices := range nls {
		mutations.PushBack(Mutation{nlIndices[0], nlIndices[1], "
\n"})
	}

	mutList := make(MutationList, mutations.Len())
	for i, el := 0, mutations.Front(); el != nil; i, el = i+1, el.Next() {
		mutList[i] = el.Value.(Mutation)
	}

	return mutList
}

func mutateTweetText(data map[string]interface{}) string {
	text := data["text"].(string)
	mutations := makeTweetMutations(data)
	sort.Sort(mutations)

	var buf bytes.Buffer
	i := 0
	for _, mutation := range mutations {
		if i < mutation.Start {
			buf.WriteString(text[i:mutation.Start])
		}
		buf.WriteString(mutation.Html)
		i = mutation.End
	}
	// Include any trailing plain text.
	buf.WriteString(text[i:])

	return buf.String()
}

func ImportJson(path string) {
	logr.Debugln("Importing from Twitter export", path)
	jsons, err := ioutil.ReadDir(path)
	if err != nil {
		logr.Errln("Error finding Twitter export", path, "to import:", err.Error())
		return
	}

	count := 0
	for _, fileinfo := range jsons {
		if fileinfo.IsDir() {
			continue
		}
		if !strings.HasSuffix(fileinfo.Name(), "json") {
			continue
		}

		datafilepath := filepath.Join(path, fileinfo.Name())
		datafile, err := os.Open(datafilepath)
		if err != nil {
			logr.Errln("Error opening Twitter export file", datafilepath, ":", err.Error())
			return
		}

		var data map[string]interface{}
		dec := json.NewDecoder(datafile)
		err = dec.Decode(&data)
		if err != nil {
			logr.Errln("Error unmarshaling Twitter export file", datafilepath, ":", err.Error())
			return
		}

		if replyId, ok := data["in_reply_to_status_id_str"]; ok && replyId != nil && replyId.(string) != "" {
			continue
		}
		if retweeted, ok := data["retweeted"]; ok && retweeted != nil && retweeted.(bool) {
			continue
		}

		tweetId := data["id_str"].(string)
		im, err := ImportBySourceIdentifier("twitter", tweetId)
		if err == sql.ErrNoRows {
			im = NewImport()
			im.Source = "twitter"
			im.Identifier = tweetId
		} else if err != nil {
			logr.Errln("Error searching for existing imported post (twitter,", tweetId, "):", err.Error())
			return
		}

		var post *Post
		if im.PostId != 0 {
			post, err = PostById(im.PostId)
			if err != nil {
				logr.Errln("Error loading already-imported post", im.PostId, "for twitter post", im.Identifier, ":", err.Error())
				return
			}
		} else {
			post = NewPost()
		}

		tweetDate := data["created_at"].(string)
		post.Posted, err = time.Parse(time.RubyDate, tweetDate)
		if err != nil {
			logr.Errln("Error parsing publish time", tweetDate, "for twitter post", tweetId, ":", err.Error())
			return
		}

		post.Html = mutateTweetText(data)

		// TODO: store the source?
		// TODO: store the geoplace

		err = post.Save()
		if err != nil {
			logr.Errln("Error saving imported post:", err.Error())
			return
		}

		im.PostId = post.Id
		err = im.Save()
		if err != nil {
			logr.Errln("Error saving import notation for post", im.PostId, ":", err.Error())
			return
		}

		logr.Debugln("Imported post (twitter,", im.Identifier, ")")
		count++
	}

	logr.Debugln("Imported", count, "posts")
}

func ImportThinkup(path string) {
	logr.Debugln("Importing from Thinkup export", path)
	port, err := os.Open(path)
	if err != nil {
		logr.Errln("Error opening", path, "for import:", err.Error())
		return
	}

	r := csv.NewReader(port)
	// There may be missing header columns, so turn off field count checking.
	r.FieldsPerRecord = -1

	head, err := r.Read()
	if err != nil {
		logr.Errln("Error reading from import file", path, ":", err.Error())
		return
	}

	count := 0
	for {
		record, err := r.Read()
		if err != nil {
			break
		}

		data := make(map[string]string)
		for i, field := range head {
			data[field] = record[i]
		}

		// TODO: import replies, once there's something reasonable to import them as.
		if data["in_reply_to_post_id"] != "" {
			logr.Debugln("Skipping post (twitter,", data["post_id"], ") as it is a reply")
			continue
		}
		// TODO: import repeats, once there's something reasonable to import them as.
		if data["in_retweet_of_post_id"] != "" {
			logr.Debugln("Skipping post (twitter,", data["post_id"], ") as it is a repeat")
			continue
		}

		// okay now what
		im, err := ImportBySourceIdentifier("twitter", data["post_id"])
		if err == sql.ErrNoRows {
			im = NewImport()
			im.Source = "twitter"
			im.Identifier = data["post_id"]
		} else if err != nil {
			logr.Errln("Error searching for existing imported post (twitter,", data["post_id"], "):", err.Error())
			return
		}

		var post *Post
		if im.PostId != 0 {
			post, err = PostById(im.PostId)
			if err != nil {
				logr.Errln("Error loading already-imported post", im.PostId, "for twitter post", im.Identifier, ":", err.Error())
				return
			}
		} else {
			post = NewPost()
		}

		post.Posted, err = time.Parse("2006-01-02 15:04:05", data["pub_date"])
		if err != nil {
			logr.Errln("Error parsing publish time", data["pub_date"], "for twitter post", data["post_id"], ":", err.Error())
			return
		}

		// TODO: make the links link.
		html := template.HTMLEscapeString(data["post_text"])
		html = strings.Replace(html, "\n", "
\n", -1)
		post.Html = html

		// TODO: store the source?
		// TODO: store the geoplace

		err = post.Save()
		if err != nil {
			logr.Errln("Error saving imported post:", err.Error())
			return
		}

		im.PostId = post.Id
		err = im.Save()
		if err != nil {
			logr.Errln("Error saving import notation for post", im.PostId, ":", err.Error())
			return
		}

		logr.Debugln("Imported post (twitter,", im.Identifier, ")")
		count++
	}
	if err != nil {
		logr.Errln("Error reading import records:", err.Error())
		return
	}

	logr.Debugln("Finished importing", count, "posts!")
}
<commit_msg>Mutation indices are by rune, not byte<commit_after>package main

import (
	"bytes"
	"container/list"
	"database/sql"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"html"
	"html/template"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"time"
)

type Import struct {
	Id         uint64 `col:"id"`
	PostId     uint64 `col:"post"`
	Source     string `col:"source"`
	Identifier string `col:"identifier"`
}

func NewImport() *Import {
	return &Import{0, 0, "", ""}
}

func ImportBySourceIdentifier(source, identifier string) (*Import, error) {
	row := db.QueryRow("SELECT id, post, source, identifier FROM import WHERE source = $1 AND identifier = $2 LIMIT 1",
		source, identifier)

	var id, postid uint64
	err := row.Scan(&id, &postid, &source, &identifier)
	if err != nil {
		return nil, err
	}

	i := &Import{id, postid, source, identifier}
	return i, nil
}

func (im *Import) Save() error {
	return db.Save(im, "import")
}

type Mutation struct {
	Start int
	End   int
	Html  string
}

type MutationList []Mutation

func (ml MutationList) Len() int {
	return len(ml)
}

func (ml MutationList) Less(i, j int) bool {
	mutI := ml[i]
	mutJ := ml[j]
	if mutI.Start < mutJ.Start {
		return true
	}
	if mutI.Start == mutJ.Start && mutI.End < mutJ.End {
		return true
	}
	return false
}

func (ml MutationList) Swap(i, j int) {
	ml[i], ml[j] = ml[j], ml[i]
}

func indicesForEntity(ent map[string]interface{}) (int, int) {
	indices := ent["indices"].([]interface{})
	s, e := indices[0].(float64), indices[1].(float64)
	return int(s), int(e)
}

func makeTweetMutations(data map[string]interface{}) MutationList {
	text := data["text"].(string)
	ents := data["entities"].(map[string]interface{})

	mutations := list.New()
	for _, entIf := range ents["user_mentions"].([]interface{}) {
		ent := entIf.(map[string]interface{})
		screenName := html.EscapeString(ent["screen_name"].(string))
		html := fmt.Sprintf(`<a href="https://twitter.com/%s" title="%s">@%s</a>`,
			screenName, html.EscapeString(ent["name"].(string)), screenName)
		start, end := indicesForEntity(ent)
		mutations.PushBack(Mutation{start, end, html})
	}
	for _, entIf := range ents["hashtags"].([]interface{}) {
		ent := entIf.(map[string]interface{})
		tagText := ent["text"].(string)
		html := fmt.Sprintf(`<a href="https://twitter.com/search?q=%%23%s">#%s</a>`,
			tagText, tagText)
		start, end := indicesForEntity(ent)
		mutations.PushBack(Mutation{start, end, html})
	}
	for _, urlEnts := range []interface{}{ents["urls"], ents["media"]} {
		if urlEnts == nil {
			continue
		}
		for _, entIf := range urlEnts.([]interface{}) {
			ent := entIf.(map[string]interface{})
			url := ent["expanded_url"].(string)
			if url == "" {
				url = ent["url"].(string)
			}
			linkText := ent["display_url"].(string)
			html := fmt.Sprintf(`<a href="%s">%s</a>`, url, linkText)
			start, end := indicesForEntity(ent)
			mutations.PushBack(Mutation{start, end, html})
		}
	}

	// We don't strictly need to regexp this of course but the strings package
	// won't find *all* instances of a substring, only the first or last.
	// Sadface that we can't just use lookahead assertion too.
	ampsRE, _ := regexp.Compile(`&`)
	amps := ampsRE.FindAllStringIndex(text, -1)
	for _, ampIndices := range amps {
		rest := text[ampIndices[1]:]
		matched, _ := regexp.MatchString("^(?:lt|gt|amp);", rest)
		if !matched {
			mutations.PushBack(Mutation{ampIndices[0], ampIndices[1], "&"})
		}
	}

	nlRE, _ := regexp.Compile(`\n`)
	nls := nlRE.FindAllStringIndex(text, -1)
	for _, nlIndices := range nls {
		mutations.PushBack(Mutation{nlIndices[0], nlIndices[1], "
\n"})
	}

	mutList := make(MutationList, mutations.Len())
	for i, el := 0, mutations.Front(); el != nil; i, el = i+1, el.Next() {
		mutList[i] = el.Value.(Mutation)
	}

	return mutList
}

func mutateTweetText(data map[string]interface{}) string {
	mutations := makeTweetMutations(data)
	sort.Sort(mutations)

	var buf bytes.Buffer
	text := []rune(data["text"].(string))
	i := 0
	for _, mutation := range mutations {
		if i < mutation.Start {
			buf.WriteString(string(text[i:mutation.Start]))
		}
		buf.WriteString(mutation.Html)
		i = mutation.End
	}
	// Include any trailing plain text.
	buf.WriteString(string(text[i:]))

	return buf.String()
}

func ImportJson(path string) {
	logr.Debugln("Importing from Twitter export", path)
	jsons, err := ioutil.ReadDir(path)
	if err != nil {
		logr.Errln("Error finding Twitter export", path, "to import:", err.Error())
		return
	}

	count := 0
	for _, fileinfo := range jsons {
		if fileinfo.IsDir() {
			continue
		}
		if !strings.HasSuffix(fileinfo.Name(), "json") {
			continue
		}

		datafilepath := filepath.Join(path, fileinfo.Name())
		datafile, err := os.Open(datafilepath)
		if err != nil {
			logr.Errln("Error opening Twitter export file", datafilepath, ":", err.Error())
			return
		}

		var data map[string]interface{}
		dec := json.NewDecoder(datafile)
		err = dec.Decode(&data)
		if err != nil {
			logr.Errln("Error unmarshaling Twitter export file", datafilepath, ":", err.Error())
			return
		}

		if replyId, ok := data["in_reply_to_status_id_str"]; ok && replyId != nil && replyId.(string) != "" {
			continue
		}
		if retweeted, ok := data["retweeted"]; ok && retweeted != nil && retweeted.(bool) {
			continue
		}

		tweetId := data["id_str"].(string)
		im, err := ImportBySourceIdentifier("twitter", tweetId)
		if err == sql.ErrNoRows {
			im = NewImport()
			im.Source = "twitter"
			im.Identifier = tweetId
		} else if err != nil {
			logr.Errln("Error searching for existing imported post (twitter,", tweetId, "):", err.Error())
			return
		}

		var post *Post
		if im.PostId != 0 {
			post, err = PostById(im.PostId)
			if err != nil {
				logr.Errln("Error loading already-imported post", im.PostId, "for twitter post", im.Identifier, ":", err.Error())
				return
			}
		} else {
			post = NewPost()
		}

		tweetDate := data["created_at"].(string)
		post.Posted, err = time.Parse(time.RubyDate, tweetDate)
		if err != nil {
			logr.Errln("Error parsing publish time", tweetDate, "for twitter post", tweetId, ":", err.Error())
			return
		}

		post.Html = mutateTweetText(data)

		// TODO: store the source?
		// TODO: store the geoplace

		err = post.Save()
		if err != nil {
			logr.Errln("Error saving imported post:", err.Error())
			return
		}

		im.PostId = post.Id
		err = im.Save()
		if err != nil {
			logr.Errln("Error saving import notation for post", im.PostId, ":", err.Error())
			return
		}

		logr.Debugln("Imported post (twitter,", im.Identifier, ")")
		count++
	}

	logr.Debugln("Imported", count, "posts")
}

func ImportThinkup(path string) {
	logr.Debugln("Importing from Thinkup export", path)
	port, err := os.Open(path)
	if err != nil {
		logr.Errln("Error opening", path, "for import:", err.Error())
		return
	}

	r := csv.NewReader(port)
	// There may be missing header columns, so turn off field count checking.
	r.FieldsPerRecord = -1

	head, err := r.Read()
	if err != nil {
		logr.Errln("Error reading from import file", path, ":", err.Error())
		return
	}

	count := 0
	for {
		record, err := r.Read()
		if err != nil {
			break
		}

		data := make(map[string]string)
		for i, field := range head {
			data[field] = record[i]
		}

		// TODO: import replies, once there's something reasonable to import them as.
		if data["in_reply_to_post_id"] != "" {
			logr.Debugln("Skipping post (twitter,", data["post_id"], ") as it is a reply")
			continue
		}
		// TODO: import repeats, once there's something reasonable to import them as.
		if data["in_retweet_of_post_id"] != "" {
			logr.Debugln("Skipping post (twitter,", data["post_id"], ") as it is a repeat")
			continue
		}

		// okay now what
		im, err := ImportBySourceIdentifier("twitter", data["post_id"])
		if err == sql.ErrNoRows {
			im = NewImport()
			im.Source = "twitter"
			im.Identifier = data["post_id"]
		} else if err != nil {
			logr.Errln("Error searching for existing imported post (twitter,", data["post_id"], "):", err.Error())
			return
		}

		var post *Post
		if im.PostId != 0 {
			post, err = PostById(im.PostId)
			if err != nil {
				logr.Errln("Error loading already-imported post", im.PostId, "for twitter post", im.Identifier, ":", err.Error())
				return
			}
		} else {
			post = NewPost()
		}

		post.Posted, err = time.Parse("2006-01-02 15:04:05", data["pub_date"])
		if err != nil {
			logr.Errln("Error parsing publish time", data["pub_date"], "for twitter post", data["post_id"], ":", err.Error())
			return
		}

		// TODO: make the links link.
		html := template.HTMLEscapeString(data["post_text"])
		html = strings.Replace(html, "\n", "
\n", -1)
		post.Html = html

		// TODO: store the source?
		// TODO: store the geoplace

		err = post.Save()
		if err != nil {
			logr.Errln("Error saving imported post:", err.Error())
			return
		}

		im.PostId = post.Id
		err = im.Save()
		if err != nil {
			logr.Errln("Error saving import notation for post", im.PostId, ":", err.Error())
			return
		}

		logr.Debugln("Imported post (twitter,", im.Identifier, ")")
		count++
	}
	if err != nil {
		logr.Errln("Error reading import records:", err.Error())
		return
	}

	logr.Debugln("Finished importing", count, "posts!")
}
<|endoftext|>